TREE3D

tree3d.py

import math
import random
import traceback

from pyrx import Ap, Db, Ge

# ---------------- CONFIG ----------------
TOTAL_HEIGHT = 200.0
BASE_RADIUS = 60.0
TIERS = 4

TRUNK_HEIGHT = 25.0
TRUNK_RADIUS = 8.0

BULBS = 120
BULB_RADIUS = 2.5
SEED = 42

BULB_COLORS = [1, 2, 3, 5, 6, 7]
FOLIAGE_COLOR = 3
TRUNK_COLOR = 53
STAR_COLOR = 2  # gold-ish yellow
STAR_RADIUS = 6.0  # size of central star sphere
STAR_CONE_HEIGHT = 12.0  # length of star rays
STAR_CONE_RADIUS = 3.0

# ----------------------------------------
def point_on_cone(base_r, top_r, base_z, top_z, u, theta):
    z = base_z + u * (top_z - base_z)
    r = base_r + u * (top_r - base_r)
    x = r * math.cos(theta)
    y = r * math.sin(theta)
    return Ge.Point3d(x, y, z)

def build_tiers():
    tiers = []
    t_h = TOTAL_HEIGHT / TIERS
    for i in range(TIERS):
        base_z = TRUNK_HEIGHT + i * t_h
        top_z = base_z + t_h
        f1 = 1.0 - (i / TIERS)
        f2 = 1.0 - ((i + 1) / TIERS)
        tiers.append(
            {
                "base_z": base_z,
                "top_z": top_z,
                "base_r": BASE_RADIUS * f1,
                "top_r": BASE_RADIUS * f2,
            }
        )
    return tiers


# ---------------- STAR CREATION ----------------
def add_star(ms: Db.BlockTableRecord, apex_z):
    """
    Creates a glowing 3D star made of:
    - A central sphere
    - 5 horizontal radial cones (around axis)
    - 1 vertical upward cone
    """

    # ---- Central sphere ----
    core = Db.Solid3d()
    core.createSphere(STAR_RADIUS)
    m = Ge.Matrix3d.translation(Ge.Vector3d(0, 0, apex_z + STAR_RADIUS))
    core.transformBy(m)
    core.setColorIndex(STAR_COLOR)
    ms.appendAcDbEntity(core)

    # ---- Five side rays ----
    for i in range(5):
        ang = i * (2 * math.pi / 5)

        cone = Db.Solid3d()
        cone.createCone(STAR_CONE_HEIGHT, STAR_CONE_RADIUS, 0.1)

        # orient cone horizontally
        axis = Ge.Vector3d(math.cos(ang), math.sin(ang), 0)
        rot = Ge.Matrix3d.rotation(math.pi / 2, Ge.Vector3d(0, 1, 0),Ge.Point3d.kOrigin)
        cone.transformBy(rot)

        # rotate into correct direction
        zaxis = Ge.Vector3d(0, 0, 1)
        rot2 = Ge.Matrix3d.rotation(math.atan2(axis.y, axis.x), zaxis,Ge.Point3d.kOrigin)
        cone.transformBy(rot2)

        # move to apex
        disp = Ge.Matrix3d.translation(Ge.Vector3d(0, 0, apex_z + STAR_RADIUS))
        cone.transformBy(disp)

        cone.setColorIndex(STAR_COLOR)
        ms.appendAcDbEntity(cone)

    # ---- Upward vertical ray ----
    up = Db.Solid3d()
    up.createCone(STAR_CONE_HEIGHT, STAR_CONE_RADIUS, 0.1)

    # point upward: cone already points up, just move
    disp2 = Ge.Matrix3d.translation(Ge.Vector3d(0, 0, apex_z + STAR_RADIUS + STAR_CONE_HEIGHT))
    up.transformBy(disp2)
    up.setColorIndex(STAR_COLOR)
    ms.appendAcDbEntity(up)


# ---------------- MAIN COMMAND ----------------
@Ap.Command("TREE3D")
def cmdTREE3D():
    try:
        if SEED is not None:
            random.seed(SEED)
        tiers = build_tiers()
        db = Db.curDb()
        ms = db.modelSpace(Db.OpenMode.kForWrite)

        # ---- TRUNK ----
        trunk = Db.Solid3d()
        trunk.createCylinder(TRUNK_HEIGHT, TRUNK_RADIUS)
        move = Ge.Matrix3d.translation(Ge.Vector3d(0, 0, TRUNK_HEIGHT / 2))
        trunk.transformBy(move)
        trunk.setColorIndex(TRUNK_COLOR)
        ms.appendAcDbEntity(trunk)

        # ---- FOLIAGE ----
        for t in tiers:
            h = t["top_z"] - t["base_z"]
            cone = Db.Solid3d()
            cone.createCone(h, t["base_r"], t["top_r"])

            move = Ge.Matrix3d.translation(Ge.Vector3d(0, 0, t["base_z"] + h / 2))
            cone.transformBy(move)
            cone.setColorIndex(FOLIAGE_COLOR)
            ms.appendAcDbEntity(cone)

        # ---- BULBS ----
        # area-proportional distribution
        total_area = 0.0
        areas = []
        for t in tiers:
            h = t["top_z"] - t["base_z"]
            s = math.hypot(t["base_r"] - t["top_r"], h)
            a = math.pi * (t["base_r"] + t["top_r"]) * s
            areas.append(a)
            total_area += a

        for idx, t in enumerate(tiers):
            h = t["top_z"] - t["base_z"]
            bulbs_here = max(3, int(BULBS * (areas[idx] / total_area)))

            for _ in range(bulbs_here):
                u = random.uniform(0.05, 0.95)
                theta = random.uniform(0, 2 * math.pi)
                pos = point_on_cone(t["base_r"], t["top_r"], t["base_z"], t["top_z"], u, theta)

                bulb = Db.Solid3d()
                bulb.createSphere(BULB_RADIUS)
                m = Ge.Matrix3d.translation(Ge.Vector3d(pos.x, pos.y, pos.z))
                bulb.transformBy(m)
                bulb.setColorIndex(random.choice(BULB_COLORS))
                ms.appendAcDbEntity(bulb)

        # ---- STAR ----
        apex_z = tiers[-1]["top_z"]
        add_star(ms, apex_z)

        print("\n3D Christmas tree with glowing star created.\n")
    except Exception as err:
        traceback.print_exception(err)