voxcity 1.0.2__py3-none-any.whl → 1.0.13__py3-none-any.whl

voxcity/simulator_gpu/core.py

@@ -0,0 +1,322 @@
+"""
+Shared core utilities for simulator_gpu.
+
+Vector and ray utilities using Taichi for GPU acceleration.
+Based on ray-tracing-one-weekend-taichi patterns.
+
+GPU Optimization Notes:
+- All functions use @ti.func for GPU inlining
+- Branchless operations preferred where possible
+- Memory coalescing friendly access patterns
+- done-flag pattern for early termination (reduces warp divergence)
+"""
+
+import taichi as ti
+import math
+
+# Type aliases for clarity
+Vector3 = ti.math.vec3
+Point3 = ti.math.vec3
+Color3 = ti.math.vec3
+
+# Constants - using Taichi static for compile-time optimization
+PI = math.pi
+TWO_PI = 2.0 * math.pi
+HALF_PI = math.pi / 2.0
+DEG_TO_RAD = math.pi / 180.0
+RAD_TO_DEG = 180.0 / math.pi
+
+# Solar constant (W/m^2) - matching PALM's solar_constant
+SOLAR_CONSTANT = 1361.0
+
+# Default extinction coefficient for vegetation
+# PALM: ext_coef = 0.6_wp (radiation_model_mod.f90 line ~890)
+EXT_COEF = 0.6
+
+# Minimum stable cosine of zenith angle
+MIN_STABLE_COSZEN = 0.0262
+
+# GPU block size hint for optimal thread occupancy
+GPU_BLOCK_SIZE = 256
+
+
+@ti.func
+def normalize(v: Vector3) -> Vector3:
+    """Normalize a vector."""
+    return v / v.norm()
+
+
+@ti.func
+def normalize_safe(v: Vector3) -> Vector3:
+    """Normalize a vector with safety check for zero-length."""
+    len_sq = v.dot(v)
+    if len_sq > 1e-10:
+        return v / ti.sqrt(len_sq)
+    return Vector3(0.0, 0.0, 1.0)
+
+
+@ti.func
+def dot(v1: Vector3, v2: Vector3) -> ti.f32:
+    """Dot product of two vectors."""
+    return v1.dot(v2)
+
+
+@ti.func
+def cross(v1: Vector3, v2: Vector3) -> Vector3:
+    """Cross product of two vectors."""
+    return v1.cross(v2)
+
+
+@ti.func
+def reflect(v: Vector3, n: Vector3) -> Vector3:
+    """Reflect vector v around normal n."""
+    return v - 2.0 * v.dot(n) * n
+
+
+@ti.func
+def ray_at(origin: Point3, direction: Vector3, t: ti.f32) -> Point3:
+    """Get point along ray at parameter t."""
+    return origin + t * direction
+
+
+@ti.func
+def length_squared(v: Vector3) -> ti.f32:
+    """Compute squared length of vector (avoids sqrt)."""
+    return v.dot(v)
+
+
+@ti.func
+def distance_squared(p1: Point3, p2: Point3) -> ti.f32:
+    """Compute squared distance between two points (avoids sqrt)."""
+    diff = p2 - p1
+    return diff.dot(diff)
+
+
+@ti.func
+def min3(a: ti.f32, b: ti.f32, c: ti.f32) -> ti.f32:
+    """Branchless minimum of three values."""
+    return ti.min(a, ti.min(b, c))
+
+
+@ti.func
+def max3(a: ti.f32, b: ti.f32, c: ti.f32) -> ti.f32:
+    """Branchless maximum of three values."""
+    return ti.max(a, ti.max(b, c))
+
+
+@ti.func
+def clamp(x: ti.f32, lo: ti.f32, hi: ti.f32) -> ti.f32:
+    """Clamp value to range [lo, hi]."""
+    return ti.max(lo, ti.min(hi, x))
+
+
+@ti.func
+def random_in_unit_sphere() -> Vector3:
+    """Generate random point in unit sphere."""
+    theta = ti.random() * TWO_PI
+    v = ti.random()
+    phi = ti.acos(2.0 * v - 1.0)
+    r = ti.random() ** (1.0 / 3.0)
+    return Vector3(
+        r * ti.sin(phi) * ti.cos(theta),
+        r * ti.sin(phi) * ti.sin(theta),
+        r * ti.cos(phi)
+    )
+
+
+@ti.func
+def random_in_hemisphere(normal: Vector3) -> Vector3:
+    """Generate random vector in hemisphere around normal."""
+    vec = random_in_unit_sphere()
+    if vec.dot(normal) < 0.0:
+        vec = -vec
+    return vec
+
+
+@ti.func
+def random_cosine_hemisphere(normal: Vector3) -> Vector3:
+    """
+    Generate random vector with cosine-weighted distribution in hemisphere.
+    Used for diffuse radiation sampling.
+    """
+    u1 = ti.random()
+    u2 = ti.random()
+    r = ti.sqrt(u1)
+    theta = TWO_PI * u2
+    
+    x = r * ti.cos(theta)
+    y = r * ti.sin(theta)
+    z = ti.sqrt(1.0 - u1)
+    
+    # Create orthonormal basis around normal
+    up = Vector3(0.0, 1.0, 0.0)
+    if ti.abs(normal.y) > 0.999:
+        up = Vector3(1.0, 0.0, 0.0)
+    
+    tangent = normalize(cross(up, normal))
+    bitangent = cross(normal, tangent)
+    
+    return normalize(x * tangent + y * bitangent + z * normal)
+
+
+@ti.func 
+def spherical_to_cartesian(azimuth: ti.f32, elevation: ti.f32) -> Vector3:
+    """
+    Convert spherical coordinates to Cartesian unit vector.
+    
+    Args:
+        azimuth: Angle from north (y-axis), clockwise, in radians
+        elevation: Angle from horizontal, in radians (0 = horizontal, pi/2 = zenith)
+    
+    Returns:
+        Unit vector (x, y, z) where z is vertical (up)
+    """
+    cos_elev = ti.cos(elevation)
+    sin_elev = ti.sin(elevation)
+    cos_azim = ti.cos(azimuth)
+    sin_azim = ti.sin(azimuth)
+    
+    # x = east, y = north, z = up
+    x = cos_elev * sin_azim
+    y = cos_elev * cos_azim
+    z = sin_elev
+    
+    return Vector3(x, y, z)
+
+
+@ti.func
+def cartesian_to_spherical(v: Vector3) -> ti.math.vec2:
+    """
+    Convert Cartesian unit vector to spherical coordinates.
+    
+    Returns:
+        vec2(azimuth, elevation) in radians
+    """
+    elevation = ti.asin(ti.math.clamp(v.z, -1.0, 1.0))
+    azimuth = ti.atan2(v.x, v.y)
+    if azimuth < 0.0:
+        azimuth += TWO_PI
+    return ti.math.vec2(azimuth, elevation)
+
+
+@ti.func
+def rotate_vector_axis_angle(vec: Vector3, axis: Vector3, angle: ti.f32) -> Vector3:
+    """
+    Rotate vector around an axis by a given angle using Rodrigues' rotation formula.
+    
+    Args:
+        vec: Vector to rotate
+        axis: Rotation axis (will be normalized)
+        angle: Rotation angle in radians
+    
+    Returns:
+        Rotated vector
+    """
+    axis_len = axis.norm()
+    if axis_len < 1e-12:
+        return vec
+    
+    k = axis / axis_len
+    c = ti.cos(angle)
+    s = ti.sin(angle)
+    
+    # Rodrigues' rotation formula: v_rot = v*cos(θ) + (k×v)*sin(θ) + k*(k·v)*(1-cos(θ))
+    v_rot = vec * c + cross(k, vec) * s + k * dot(k, vec) * (1.0 - c)
+    
+    return v_rot
+
+
+@ti.func
+def build_face_basis(normal: Vector3):
+    """
+    Build orthonormal basis for a face with given normal.
+    
+    Returns:
+        Tuple of (tangent, bitangent, normal) vectors
+    """
+    n_len = normal.norm()
+    if n_len < 1e-12:
+        return Vector3(1.0, 0.0, 0.0), Vector3(0.0, 1.0, 0.0), Vector3(0.0, 0.0, 1.0)
+    
+    n = normal / n_len
+    
+    # Choose helper vector not parallel to normal
+    helper = Vector3(0.0, 0.0, 1.0)
+    if ti.abs(n.z) > 0.999:
+        helper = Vector3(1.0, 0.0, 0.0)
+    
+    # Compute tangent via cross product
+    u = cross(helper, n)
+    u_len = u.norm()
+    if u_len < 1e-12:
+        u = Vector3(1.0, 0.0, 0.0)
+    else:
+        u = u / u_len
+    
+    # Bitangent
+    v = cross(n, u)
+    
+    return u, v, n
+
+
+@ti.data_oriented
+class Rays:
+    """
+    Array of rays for batch processing.
+    Similar to ray-tracing-one-weekend-taichi but adapted for view/solar tracing.
+    """
+    
+    def __init__(self, n_rays: int):
+        self.n_rays = n_rays
+        self.origin = ti.Vector.field(3, dtype=ti.f32, shape=(n_rays,))
+        self.direction = ti.Vector.field(3, dtype=ti.f32, shape=(n_rays,))
+        self.transparency = ti.field(dtype=ti.f32, shape=(n_rays,))
+        self.active = ti.field(dtype=ti.i32, shape=(n_rays,))
+    
+    @ti.func
+    def set(self, idx: ti.i32, origin: Point3, direction: Vector3, transp: ti.f32):
+        self.origin[idx] = origin
+        self.direction[idx] = direction
+        self.transparency[idx] = transp
+        self.active[idx] = 1
+    
+    @ti.func
+    def get(self, idx: ti.i32):
+        return self.origin[idx], self.direction[idx], self.transparency[idx]
+    
+    @ti.func
+    def deactivate(self, idx: ti.i32):
+        self.active[idx] = 0
+    
+    @ti.func
+    def is_active(self, idx: ti.i32) -> ti.i32:
+        return self.active[idx]
+
+
+@ti.data_oriented
+class HitRecord:
+    """
+    Store ray-surface intersection results.
+    """
+    
+    def __init__(self, n_rays: int):
+        self.n_rays = n_rays
+        self.hit = ti.field(dtype=ti.i32, shape=(n_rays,))
+        self.t = ti.field(dtype=ti.f32, shape=(n_rays,))
+        self.point = ti.Vector.field(3, dtype=ti.f32, shape=(n_rays,))
+        self.normal = ti.Vector.field(3, dtype=ti.f32, shape=(n_rays,))
+        self.surface_id = ti.field(dtype=ti.i32, shape=(n_rays,))
+    
+    @ti.func
+    def set(self, idx: ti.i32, hit: ti.i32, t: ti.f32, point: Point3, 
+            normal: Vector3, surface_id: ti.i32):
+        self.hit[idx] = hit
+        self.t[idx] = t
+        self.point[idx] = point
+        self.normal[idx] = normal
+        self.surface_id[idx] = surface_id
+    
+    @ti.func
+    def get(self, idx: ti.i32):
+        return (self.hit[idx], self.t[idx], self.point[idx], 
+                self.normal[idx], self.surface_id[idx])

voxcity/simulator_gpu/domain.py

@@ -0,0 +1,262 @@
+"""
+Shared domain definition for simulator_gpu.
+
+Represents the 3D computational domain with:
+- Grid cells (dx, dy, dz spacing)
+- Topography (terrain height)
+- Building geometry (3D obstacles)
+- Plant canopy (Leaf Area Density - LAD)
+- Tree mask for view analysis
+"""
+
+import taichi as ti
+import numpy as np
+from typing import Tuple, Optional, Union
+from .core import Vector3, Point3, EXT_COEF
+from .init_taichi import ensure_initialized
+
+
+# Surface direction indices (matching PALM convention)
+IUP = 0      # Upward facing (horizontal roof/ground)
+IDOWN = 1    # Downward facing
+INORTH = 2   # North facing (positive y)
+ISOUTH = 3   # South facing (negative y)
+IEAST = 4    # East facing (positive x)
+IWEST = 5    # West facing (negative x)
+
+# Direction normal vectors (x, y, z)
+DIR_NORMALS = {
+    IUP: (0.0, 0.0, 1.0),
+    IDOWN: (0.0, 0.0, -1.0),
+    INORTH: (0.0, 1.0, 0.0),
+    ISOUTH: (0.0, -1.0, 0.0),
+    IEAST: (1.0, 0.0, 0.0),
+    IWEST: (-1.0, 0.0, 0.0),
+}
+
+
+@ti.data_oriented
+class Domain:
+    """
+    3D computational domain for simulation.
+    
+    The domain uses a regular grid with:
+    - x: West to East
+    - y: South to North  
+    - z: Ground to Sky
+    
+    Attributes:
+        nx, ny, nz: Number of grid cells in each direction
+        dx, dy, dz: Grid spacing in meters
+        origin: (x, y, z) coordinates of domain origin
+    """
+    
+    def __init__(
+        self,
+        nx: int,
+        ny: int, 
+        nz: int,
+        dx: float = 1.0,
+        dy: float = 1.0,
+        dz: float = 1.0,
+        origin: Tuple[float, float, float] = (0.0, 0.0, 0.0),
+        origin_lat: Optional[float] = None,
+        origin_lon: Optional[float] = None
+    ):
+        """
+        Initialize the domain.
+        
+        Args:
+            nx, ny, nz: Grid dimensions
+            dx, dy, dz: Grid spacing (m)
+            origin: Domain origin coordinates
+            origin_lat: Latitude for solar calculations (degrees)
+            origin_lon: Longitude for solar calculations (degrees)
+        """
+        # Ensure Taichi is initialized before creating any fields
+        ensure_initialized()
+        
+        self.nx = nx
+        self.ny = ny
+        self.nz = nz
+        self.dx = dx
+        self.dy = dy
+        self.dz = dz
+        self.origin = origin
+        self.origin_lat = origin_lat if origin_lat is not None else 0.0
+        self.origin_lon = origin_lon if origin_lon is not None else 0.0
+        
+        # Domain bounds
+        self.x_min = origin[0]
+        self.x_max = origin[0] + nx * dx
+        self.y_min = origin[1]
+        self.y_max = origin[1] + ny * dy
+        self.z_min = origin[2]
+        self.z_max = origin[2] + nz * dz
+        
+        # Grid cell volume
+        self.cell_volume = dx * dy * dz
+        
+        # Topography: terrain height at each (i, j) column
+        self.topo_top = ti.field(dtype=ti.i32, shape=(nx, ny))
+        
+        # Building mask: 1 if cell is solid (building), 0 if air
+        self.is_solid = ti.field(dtype=ti.i32, shape=(nx, ny, nz))
+        
+        # Tree mask: 1 if cell is tree canopy, 0 otherwise
+        self.is_tree = ti.field(dtype=ti.i32, shape=(nx, ny, nz))
+        
+        # Leaf Area Density (m^2/m^3) for plant canopy
+        self.lad = ti.field(dtype=ti.f32, shape=(nx, ny, nz))
+        
+        # Plant canopy top index for each column
+        self.plant_top = ti.field(dtype=ti.i32, shape=(nx, ny))
+        
+        # Surface count
+        self.n_surfaces = ti.field(dtype=ti.i32, shape=())
+        
+        # Initialize arrays
+        self._init_arrays()
+        
+    @ti.kernel
+    def _init_arrays(self):
+        """Initialize all arrays to default values."""
+        for i, j in self.topo_top:
+            self.topo_top[i, j] = 0
+            self.plant_top[i, j] = 0
+        
+        for i, j, k in self.is_solid:
+            self.is_solid[i, j, k] = 0
+            self.is_tree[i, j, k] = 0
+            self.lad[i, j, k] = 0.0
+    
+    def set_flat_terrain(self, height: float = 0.0):
+        """Set flat terrain at given height."""
+        k_top = int(height / self.dz)
+        self._set_flat_terrain_kernel(k_top)
+    
+    def initialize_terrain(self, height: float = 0.0):
+        """Alias for set_flat_terrain."""
+        self.set_flat_terrain(height)
+    
+    @ti.kernel
+    def _set_flat_terrain_kernel(self, k_top: ti.i32):
+        for i, j in self.topo_top:
+            self.topo_top[i, j] = k_top
+            for k in range(k_top + 1):
+                self.is_solid[i, j, k] = 1
+    
+    def set_terrain_from_array(self, terrain_height: np.ndarray):
+        """
+        Set terrain from 2D numpy array of heights.
+        
+        Args:
+            terrain_height: 2D array (nx, ny) of terrain heights in meters
+        """
+        terrain_k = (terrain_height / self.dz).astype(np.int32)
+        self._set_terrain_kernel(terrain_k)
+    
+    @ti.kernel
+    def _set_terrain_kernel(self, terrain_k: ti.types.ndarray()):
+        for i, j in self.topo_top:
+            k_top = terrain_k[i, j]
+            self.topo_top[i, j] = k_top
+            for k in range(self.nz):
+                if k <= k_top:
+                    self.is_solid[i, j, k] = 1
+                else:
+                    self.is_solid[i, j, k] = 0
+    
+    def add_building(
+        self,
+        x_range: Optional[Tuple[int, int]] = None,
+        y_range: Optional[Tuple[int, int]] = None,
+        z_range: Optional[Tuple[int, int]] = None,
+        *,
+        x_start: Optional[int] = None,
+        x_end: Optional[int] = None,
+        y_start: Optional[int] = None,
+        y_end: Optional[int] = None,
+        height: Optional[float] = None
+    ):
+        """
+        Add a rectangular building to the domain.
+        """
+        # Handle convenience parameters
+        if x_start is not None and x_end is not None:
+            x_range = (x_start, x_end)
+        if y_start is not None and y_end is not None:
+            y_range = (y_start, y_end)
+        if height is not None and z_range is None:
+            k_top = int(height / self.dz) + 1
+            z_range = (0, k_top)
+        
+        if x_range is None or y_range is None or z_range is None:
+            raise ValueError("Must provide either range tuples or individual parameters")
+        
+        self._add_building_kernel(x_range[0], x_range[1], y_range[0], y_range[1], z_range[0], z_range[1])
+    
+    @ti.kernel
+    def _add_building_kernel(self, i_min: ti.i32, i_max: ti.i32, j_min: ti.i32, j_max: ti.i32, k_min: ti.i32, k_max: ti.i32):
+        for i, j, k in ti.ndrange((i_min, i_max), (j_min, j_max), (k_min, k_max)):
+            self.is_solid[i, j, k] = 1
+    
+    def add_tree(
+        self,
+        x_range: Tuple[int, int],
+        y_range: Tuple[int, int],
+        z_range: Tuple[int, int],
+        lad_value: float = 1.0
+    ):
+        """
+        Add a tree canopy region to the domain.
+        
+        Args:
+            x_range, y_range, z_range: Grid index ranges
+            lad_value: Leaf Area Density value
+        """
+        self._add_tree_kernel(x_range[0], x_range[1], y_range[0], y_range[1], z_range[0], z_range[1], lad_value)
+    
+    @ti.kernel
+    def _add_tree_kernel(self, i_min: ti.i32, i_max: ti.i32, j_min: ti.i32, j_max: ti.i32, k_min: ti.i32, k_max: ti.i32, lad: ti.f32):
+        for i, j, k in ti.ndrange((i_min, i_max), (j_min, j_max), (k_min, k_max)):
+            self.is_tree[i, j, k] = 1
+            self.lad[i, j, k] = lad
+    
+    def set_from_voxel_data(self, voxel_data: np.ndarray, tree_code: int = -2, solid_codes: Optional[list] = None):
+        """
+        Set domain from a 3D voxel data array.
+        
+        Args:
+            voxel_data: 3D numpy array with voxel class codes
+            tree_code: Class code for trees (default -2)
+            solid_codes: List of codes that are solid (default: all non-zero except tree_code)
+        """
+        if solid_codes is None:
+            # All non-zero codes except tree are solid
+            solid_codes = []
+        
+        self._set_from_voxel_data_kernel(voxel_data, tree_code)
+    
+    @ti.kernel
+    def _set_from_voxel_data_kernel(self, voxel_data: ti.types.ndarray(), tree_code: ti.i32):
+        for i, j, k in ti.ndrange(self.nx, self.ny, self.nz):
+            val = voxel_data[i, j, k]
+            if val == tree_code:
+                self.is_tree[i, j, k] = 1
+                self.is_solid[i, j, k] = 0
+            elif val != 0:
+                self.is_solid[i, j, k] = 1
+                self.is_tree[i, j, k] = 0
+            else:
+                self.is_solid[i, j, k] = 0
+                self.is_tree[i, j, k] = 0
+    
+    def get_max_dist(self) -> float:
+        """Get maximum ray distance (domain diagonal)."""
+        import math
+        return math.sqrt(
+            (self.nx * self.dx)**2 + 
+            (self.ny * self.dy)**2 + 
+            (self.nz * self.dz)**2
+        )

voxcity/simulator_gpu/environment.yml

@@ -0,0 +1,11 @@
+name: voxcity_gpu
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.12
+  - gdal
+  - timezonefinder
+  - pip
+  - pip:
+    - voxcity[gpu]