py-growth-RHEED 1.0.2__py3-none-any.whl

PY_GROWTH/__init__.py

@@ -0,0 +1,3 @@
+"""
+PY_GROWTH - RHEED Intensity Oscillation Simulator.
+"""

PY_GROWTH/jax/__init__.py

@@ -0,0 +1,8 @@
+from .core import run_simulation
+from .ode_solver import solve_ode_adaptive
+from .models import derivs0, derivs1, derivs2, compute_intensity, compute_rms_roughness
+
+# Re-export IO utilities from numpy module to maintain strict subpackage structural symmetry
+from PY_GROWTH.numpy.io import load_input_data, save_history_to_json
+
+

PY_GROWTH/jax/core.py

@@ -0,0 +1,105 @@
+"""
+core.py
+
+Defines the centralized pipeline that binds the physics models with the ODE 
+solver to produce and return calculation history dictionaries.
+Accelerated end-to-end using JAX.
+"""
+
+from typing import Dict, Any, Literal
+import time
+import jax
+import jax.numpy as jnp
+
+# Enable 64-bit precision to match NumPy/C++ for identical math tracing
+jax.config.update("jax_enable_x64", True)
+
+from functools import partial
+
+from .ode_solver import solve_ode_adaptive
+from .models import derivs0, derivs1, derivs2, compute_intensity, compute_rms_roughness
+
+@partial(jax.jit, static_argnames=['model_type', 'num_layers', 'num_intervals'])
+def _run_sim_jitted(
+    model_type: int,
+    num_layers: int,
+    t_max: float,
+    num_intervals: int,
+    an_kn: jnp.ndarray,
+    growth_rates: jnp.ndarray
+):
+    """
+    Core JIT-compiled simulation trace.
+    """
+    if model_type == 0:
+        derivs_closure = lambda x, y: derivs0(x, y, an_kn, growth_rates)
+    elif model_type == 1:
+        derivs_closure = lambda x, y: derivs1(x, y, an_kn, growth_rates)
+    elif model_type == 2:
+        derivs_closure = lambda x, y: derivs2(x, y, an_kn, growth_rates)
+    else:
+        # Fallback to 0 if traced incorrectly, but static_argnames prevents this
+        derivs_closure = lambda x, y: derivs0(x, y, an_kn, growth_rates)
+
+    y0 = jnp.zeros(num_layers, dtype=jnp.float64)
+    
+    results = solve_ode_adaptive(
+        derivs=derivs_closure, 
+        y0=y0, 
+        t_span=(0.0, t_max), 
+        num_intervals=num_intervals
+    )
+    
+    t_history = results["t"]
+    y_history = results["y"]
+    valid_count = results["valid_count"]
+    
+    intensity = compute_intensity(y_history)
+    roughness = compute_rms_roughness(t_history, y_history, growth_rates)
+    
+    return t_history, y_history, intensity, roughness, valid_count
+
+def run_simulation(
+    model_type: Literal[0, 1, 2],
+    num_layers: int,
+    t_max: float,
+    num_intervals: int,
+    an_kn: jnp.ndarray,
+    growth_rates: jnp.ndarray
+) -> Dict[str, Any]:
+    """
+    Orchestrates the chosen ODE implementation on a zero-based coverage system.
+    Returns standard python dictionary slicing dynamic lengths appropriately.
+    """
+    import numpy as np # For output conversion
+    
+    start_time = time.perf_counter()
+    
+    an_kn_jnp = jnp.array(an_kn, dtype=jnp.float64)
+    gR_jnp = jnp.array(growth_rates, dtype=jnp.float64)
+    
+    t_hist_jnp, y_hist_jnp, int_jnp, rough_jnp, valid_count = _run_sim_jitted(
+        model_type=int(model_type),
+        num_layers=int(num_layers),
+        t_max=float(t_max),
+        num_intervals=int(num_intervals),
+        an_kn=an_kn_jnp,
+        growth_rates=gR_jnp
+    )
+    
+    # Block until ready to measure true execution time including GPU execution
+    t_hist_jnp.block_until_ready()
+    
+    end_time = time.perf_counter()
+    calc_time = end_time - start_time
+    
+    v = int(valid_count)
+    
+    return {
+        "time": np.array(t_hist_jnp[:v]),
+        "coverage": np.array(y_hist_jnp[:v, :]),
+        "intensity": np.array(int_jnp[:v]),
+        "rms_roughness": np.array(rough_jnp[:v]),
+        "model_type": model_type,
+        "execution_time_seconds": calc_time
+    }

PY_GROWTH/jax/models.py

@@ -0,0 +1,169 @@
+"""
+models.py
+
+Vectorized mathematical models translating C++ loop calculations representing 
+diffusive and distributed RHEED epitaxial growth dynamics.
+Accelerated using JAX.
+"""
+
+import jax
+import jax.numpy as jnp
+from functools import partial
+
+@jax.jit
+def _th_padded(theta: jnp.ndarray) -> jnp.ndarray:
+    """
+    Pads the coverage array for generic layer interaction without out-of-bounds errors.
+    Upper boundary condition (n<=0) treats coverage as 1.0 (fully covered substrate).
+    Lower boundary condition (n>nMax) treats coverage as 0.0 (no coverage yet).
+    """
+    return jnp.concatenate((jnp.array([1.0, 1.0]), theta, jnp.array([0.0, 0.0])))
+
+@jax.jit
+def _dn1(theta_array: jnp.ndarray) -> jnp.ndarray:
+    """Distribution function 1 for Distributed model 1."""
+    th_c = jnp.clip(theta_array, 0.0, 1.0)
+    return th_c * jnp.sqrt(1.0 - th_c)
+
+@jax.jit
+def _dn2(theta_array: jnp.ndarray) -> jnp.ndarray:
+    """Distribution function 2 for Distributed model 2."""
+    th_c = jnp.clip(theta_array, 0.0, 1.0)
+    return jnp.where(th_c < 0.5, jnp.sqrt(th_c), jnp.sqrt(1.0 - th_c))
+
+@jax.jit
+def derivs0(t: float, theta: jnp.ndarray, C: jnp.ndarray, gR: jnp.ndarray) -> jnp.ndarray:
+    """ Diffusive growth: Model 0 """
+    th = _th_padded(theta)
+    
+    th_n = th[2:-2]
+    th_n_minus_1 = th[1:-3]
+    th_n_minus_2 = th[:-4]
+    th_n_plus_1 = th[3:-1]
+    th_n_plus_2 = th[4:]
+    
+    dTheta = th_n_minus_1 - th_n
+    
+    dThetaDt = (dTheta * gR + 
+                C * (th_n_plus_1 - th_n_plus_2) * dTheta - 
+                C * (th_n - th_n_plus_1) * (th_n_minus_2 - th_n_minus_1))
+    
+    return dThetaDt
+
+@jax.jit
+def derivs1(t: float, theta: jnp.ndarray, C: jnp.ndarray, gR: jnp.ndarray) -> jnp.ndarray:
+    """ Distributed growth model 1: Model 1 """
+    th = _th_padded(theta)
+    
+    th_n = th[2:-2]
+    th_n_minus_1 = th[1:-3]
+    th_n_plus_1 = th[3:-1]
+    
+    dn1_n = _dn1(th_n)
+    dn1_n_minus_1 = _dn1(th_n_minus_1)
+    dn1_n_plus_1 = _dn1(th_n_plus_1)
+    
+    C_n = C
+    C_n_minus_1 = jnp.pad(C, (1, 0), constant_values=0.0)[:-1]
+    
+    dTheta = th_n_minus_1 - th_n
+    
+    den_m1 = dn1_n_minus_1 + dn1_n
+    term_m1 = jnp.where(dn1_n_minus_1 != 0.0, 
+                       (C_n_minus_1 * dTheta * dn1_n_minus_1) / den_m1, 
+                       0.0)
+    
+    den_p1 = dn1_n + dn1_n_plus_1
+    term_p1 = jnp.where(dn1_n > 0.0, 
+                       (C_n * (th_n - th_n_plus_1) * dn1_n) / den_p1, 
+                       0.0)
+    
+    return (dTheta - term_m1 + term_p1) * gR
+
+@jax.jit
+def derivs2(t: float, theta: jnp.ndarray, C: jnp.ndarray, gR: jnp.ndarray) -> jnp.ndarray:
+    """ Distributed growth model 2: Model 2 """
+    th = _th_padded(theta)
+    
+    th_n = th[2:-2]
+    th_n_minus_1 = th[1:-3]
+    th_n_plus_1 = th[3:-1]
+    
+    dn2_n = _dn2(th_n)
+    dn2_n_minus_1 = _dn2(th_n_minus_1)
+    dn2_n_plus_1 = _dn2(th_n_plus_1)
+    
+    C_n = C
+    C_n_minus_1 = jnp.pad(C, (1, 0), constant_values=0.0)[:-1]
+    
+    dTheta = th_n_minus_1 - th_n
+    
+    den_m1 = dn2_n_minus_1 + dn2_n
+    term_m1 = jnp.where(dn2_n_minus_1 != 0.0, 
+                       (C_n_minus_1 * dTheta * dn2_n_minus_1) / den_m1, 
+                       0.0)
+    
+    den_p1 = dn2_n + dn2_n_plus_1
+    term_p1 = jnp.where(dn2_n > 0.0, 
+                       (C_n * (th_n - th_n_plus_1) * dn2_n) / den_p1, 
+                       0.0)
+    
+    return (dTheta - term_m1 + term_p1) * gR
+
+@jax.jit
+def compute_intensity(coverage_history: jnp.ndarray) -> jnp.ndarray:
+    """
+    Computes Kinematical Diffracted Intensity.
+    
+    Args:
+        coverage_history (jnp.ndarray): History matrix of shape (num_time_steps, numLayers).
+        
+    Returns:
+        jnp.ndarray: Evaluated sequence of intensity values.
+    """
+    num_layers = coverage_history.shape[1]
+    cDI = 1.0 - coverage_history[:, 0]
+    
+    n_indices = jnp.arange(1, num_layers)
+    cos_n_pi = jnp.cos(n_indices * jnp.pi)
+    
+    cov_n = coverage_history[:, :-1]
+    cov_n_plus_1 = coverage_history[:, 1:]
+    
+    sum_term = jnp.sum((cov_n - cov_n_plus_1) * cos_n_pi, axis=1)
+    # Using python if is evaluated at trace time, which works if num_layers is static.
+    # In JAX, shapes (and thus num_layers) are always statically traced.
+    if num_layers > 1:
+        cDI += sum_term
+        
+    return cDI ** 2
+
+@jax.jit
+def compute_rms_roughness(growth_time: jnp.ndarray, coverage_history: jnp.ndarray, gR: jnp.ndarray) -> jnp.ndarray:
+    """
+    Computes the RMS Roughness over time.
+    
+    Args:
+        growth_time (jnp.ndarray): Vector of time steps elapsed.
+        coverage_history (jnp.ndarray): Shape (num_time_steps, numLayers).
+        gR (jnp.ndarray): Growth rates of the layers.
+        
+    Returns:
+        jnp.ndarray: Vector of evaluated RMS roughness configurations.
+    """
+    num_layers = coverage_history.shape[1]
+    
+    sD = (growth_time * gR[0])**2 * (1.0 - coverage_history[:, 0])
+    
+    if num_layers > 1:
+        t = growth_time[:, jnp.newaxis]
+        n_indices = jnp.arange(1, num_layers)
+        cov_n = coverage_history[:, :-1]
+        cov_n_plus_1 = coverage_history[:, 1:]
+        
+        gR_inner = gR[:num_layers-1]
+        
+        sum_term = jnp.sum((n_indices - t * gR_inner)**2 * (cov_n - cov_n_plus_1), axis=1)
+        sD += sum_term
+        
+    return jnp.sqrt(sD)

PY_GROWTH/jax/ode_solver.py

@@ -0,0 +1,139 @@
+"""
+ode_solver.py
+
+Provides a generalized, pure JAX implementation of the Cash-Karp 
+fifth-order Runge-Kutta adaptive ODE solver. Designed using 
+jax.lax.while_loop to allow end-to-end JIT compilation.
+"""
+
+import jax
+import jax.numpy as jnp
+from typing import Callable, Tuple, Dict
+
+def rkck(y: jnp.ndarray, dydx: jnp.ndarray, x: float, h: float, 
+         derivs: Callable[[float, jnp.ndarray], jnp.ndarray]) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    """Cash-Karp fifth-order Runge-Kutta step."""
+    b21 = 0.2; b31 = 3.0/40.0; b32 = 9.0/40.0; b41 = 0.3; b42 = -0.9; b43 = 1.2
+    b51 = -11.0/54.0; b52 = 2.5; b53 = -70.0/27.0; b54 = 35.0/27.0
+    b61 = 1631.0/55296.0; b62 = 175.0/512.0; b63 = 575.0/13824.0
+    b64 = 44275.0/110592.0; b65 = 253.0/4096.0
+    
+    c1 = 37.0/378.0; c3 = 250.0/621.0; c4 = 125.0/594.0; c6 = 512.0/1771.0
+    dc5 = -277.00/14336.0
+    
+    dc1 = c1 - 2825.0/27648.0; dc3 = c3 - 18575.0/48384.0
+    dc4 = c4 - 13525.0/55296.0; dc6 = c6 - 0.25
+
+    ak2 = derivs(x + b21*h, y + b21*h*dydx)
+    ak3 = derivs(x + (b31+b32)*h, y + h*(b31*dydx + b32*ak2))
+    ak4 = derivs(x + (b41+b42+b43)*h, y + h*(b41*dydx + b42*ak2 + b43*ak3))
+    ak5 = derivs(x + (b51+b52+b53+b54)*h, y + h*(b51*dydx + b52*ak2 + b53*ak3 + b54*ak4))
+    ak6 = derivs(x + (b61+b62+b63+b64+b65)*h, y + h*(b61*dydx + b62*ak2 + b63*ak3 + b64*ak4 + b65*ak5))
+
+    yout = y + h*(c1*dydx + c3*ak3 + c4*ak4 + c6*ak6)
+    yerr = h*(dc1*dydx + dc3*ak3 + dc4*ak4 + dc5*ak5 + dc6*ak6)
+    
+    return yout, yerr
+
+def rkqs(y: jnp.ndarray, dydx: jnp.ndarray, x: float, htry: float, eps: float, 
+         yscal: jnp.ndarray, derivs: Callable[[float, jnp.ndarray], jnp.ndarray]):
+    """Fifth-order Runge Kutta step with monitoring of local truncation error."""
+    SAFETY = 0.9; PGROW = -0.2; PSHRNK = -0.25; ERRCON = 1.89e-4
+    
+    def cond_fun(state):
+        h, ytemp, yerr, errmax = state
+        return errmax > 1.0
+        
+    def body_fun(state):
+        h, _, _, _ = state
+        ytemp, yerr = rkck(y, dydx, x, h, derivs)
+        errmax = jnp.max(jnp.abs(yerr / yscal)) / eps
+        
+        htemp = SAFETY * h * (errmax ** PSHRNK)
+        h_new = jnp.where(h >= 0.0, jnp.maximum(htemp, 0.1 * h), jnp.minimum(htemp, 0.1 * h))
+        
+        h_next_iter = jnp.where(errmax > 1.0, h_new, h)
+        return (h_next_iter, ytemp, yerr, errmax)
+        
+    init_errmax = jnp.array(2.0, dtype=y.dtype)
+    init_state = (htry, y, jnp.zeros_like(y), init_errmax)
+    
+    final_state = jax.lax.while_loop(cond_fun, body_fun, init_state)
+    h_used, ytemp, yerr, errmax = final_state
+    
+    hnext = jnp.where(errmax > ERRCON, 
+                      SAFETY * h_used * (errmax ** PGROW), 
+                      5.0 * h_used)
+                      
+    x_new = x + h_used
+    return ytemp, x_new, hnext, h_used
+
+def solve_ode_adaptive(derivs: Callable, y0: jnp.ndarray, t_span: Tuple[float, float], 
+                       num_intervals: int, eps: float = 1.0e-7, h1: float = 1.0e-3) -> Dict[str, jnp.ndarray]:
+    """Main driver using jax.lax.while_loop."""
+    t0, t_max = t_span
+    TINY = 1.0e-30
+    
+    d_x_sav = t_max / num_intervals
+    
+    history_t = jnp.zeros(num_intervals + 5, dtype=jnp.float64)
+    history_y = jnp.zeros((num_intervals + 5, len(y0)), dtype=jnp.float64)
+    
+    def cond_fun(state):
+        # State: y, x, h, x_sav, history_t, history_y, out_idx, done, steps
+        return jnp.logical_not(state[7])
+
+    def body_fun(state):
+        y, x, h, x_sav, history_t, history_y, out_idx, done, steps = state
+        
+        dydx = derivs(x, y)
+        yscal = jnp.abs(y) + jnp.abs(dydx * h) + TINY
+        
+        save_cond = jnp.logical_and(jnp.abs(x - x_sav) > jnp.abs(d_x_sav), out_idx < num_intervals)
+        
+        idx_to_update = jnp.where(save_cond, out_idx, 0)
+        
+        new_hist_t = jnp.where(save_cond, history_t.at[idx_to_update].set(x), history_t)
+        new_hist_y = jnp.where(save_cond, history_y.at[idx_to_update].set(y), history_y)
+        
+        new_x_sav = jnp.where(save_cond, x, x_sav)
+        new_out_idx = out_idx + jnp.where(save_cond, 1, 0)
+        
+        overshoot = (x + h - t_max) * (x + h - t0) > 0.0
+        h_adjusted = jnp.where(overshoot, t_max - x, h)
+        
+        y_next, x_next, hnext, h_used = rkqs(y, dydx, x, h_adjusted, eps, yscal, derivs)
+        
+        finished = jnp.logical_or((x_next - t_max) * (t_max - t0) >= 0.0, steps >= 300000)
+        
+        save_end_cond = jnp.logical_and(finished, new_out_idx < len(new_hist_t))
+        
+        final_idx = jnp.where(save_end_cond, new_out_idx, 0)
+        new_hist_t_final = jnp.where(save_end_cond, new_hist_t.at[final_idx].set(x_next), new_hist_t)
+        new_hist_y_final = jnp.where(save_end_cond, new_hist_y.at[final_idx].set(y_next), new_hist_y)
+        
+        final_out_idx = new_out_idx + jnp.where(save_end_cond, 1, 0)
+        
+        return (y_next, x_next, hnext, new_x_sav, new_hist_t_final, new_hist_y_final, final_out_idx, finished, steps + 1)
+
+    init_state = (
+        y0,          
+        jnp.array(t0, dtype=jnp.float64),          
+        jnp.array(h1, dtype=jnp.float64),          
+        jnp.array(t0 - d_x_sav * 2.0, dtype=jnp.float64), 
+        history_t,   
+        history_y,   
+        jnp.array(0, dtype=jnp.int32),           
+        jnp.array(False),
+        jnp.array(0, dtype=jnp.int32)
+    )
+    
+    final_state = jax.lax.while_loop(cond_fun, body_fun, init_state)
+    
+    _, _, _, _, final_history_t, final_history_y, final_out_idx, _, _ = final_state
+    
+    return {
+        "t": final_history_t,
+        "y": final_history_y,
+        "valid_count": final_out_idx
+    }

PY_GROWTH/numba/__init__.py

@@ -0,0 +1,8 @@
+from .core import run_simulation
+from .ode_solver import solve_ode_adaptive
+from .models import derivs0, derivs1, derivs2, compute_intensity, compute_rms_roughness
+
+# Re-export IO utilities from numpy module to maintain strict subpackage structural symmetry
+from PY_GROWTH.numpy.io import load_input_data, save_history_to_json
+
+

PY_GROWTH/numba/core.py

@@ -0,0 +1,49 @@
+import time
+import numpy as np
+from numba import njit
+from .ode_solver import solve_ode_adaptive
+from .models import derivs0, derivs1, derivs2, compute_intensity, compute_rms_roughness
+
+@njit(cache=True)
+def _jitted_pipeline(model_type, num_layers, t_max, num_intervals, an_kn, growth_rates):
+    """
+    Executes the entire RKCK execution block alongside physical geometry extraction mathematically isolated inside LLVM.
+    Resolves Numba limitations regarding Generic Callable python function arguments.
+    """
+    y0 = np.zeros(num_layers, dtype=np.float64)
+    
+    if model_type == 0:
+        hist_t, hist_y = solve_ode_adaptive(derivs0, y0, (0.0, t_max), (an_kn, growth_rates), num_intervals)
+    elif model_type == 1:
+        hist_t, hist_y = solve_ode_adaptive(derivs1, y0, (0.0, t_max), (an_kn, growth_rates), num_intervals)
+    else:
+        hist_t, hist_y = solve_ode_adaptive(derivs2, y0, (0.0, t_max), (an_kn, growth_rates), num_intervals)
+        
+    intensity = compute_intensity(hist_y)
+    rms = compute_rms_roughness(hist_t, hist_y, growth_rates)
+    
+    return hist_t, hist_y, intensity, rms
+
+def run_simulation(model_type: int, num_layers: int, t_max: float, num_intervals: int,
+                   an_kn: np.ndarray, growth_rates: np.ndarray) -> dict:
+    """
+    Executes the RHEED ODE simulations explicitly optimized under Numba caching mechanics.
+    """
+    if model_type not in (0, 1, 2):
+        raise ValueError(f"Unknown model_type integer: {model_type}")
+
+    start_time = time.perf_counter()
+    
+    t, y, intensity, rms = _jitted_pipeline(model_type, num_layers, t_max, num_intervals, an_kn, growth_rates)
+    
+    end_time = time.perf_counter()
+    calc_time = end_time - start_time
+    
+    return {
+        "time": t,
+        "coverage": y,
+        "intensity": intensity,
+        "rms_roughness": rms,
+        "model_type": model_type,
+        "execution_time_seconds": calc_time
+    }

PY_GROWTH/numba/models.py

@@ -0,0 +1,142 @@
+import numpy as np
+from numba import njit
+import math
+
+@njit(cache=True)
+def th(n: int, theta: np.ndarray) -> float:
+    """
+    Retrieves the coverage value simulating generalized padding.
+    Upper boundary condition (n<=0) treats coverage as 1.0 (fully covered substrate).
+    Lower boundary condition (n>nMax) treats coverage as 0.0 (no coverage yet).
+    """
+    nMax = len(theta) - 1
+    if n < 0:
+        return 1.0
+    elif n > nMax:
+        return 0.0
+    else:
+        return theta[n]
+
+@njit(cache=True)
+def dn1(n: int, theta: np.ndarray) -> float:
+    """Distribution function 1 for Distributed model 1."""
+    tt = th(n, theta)
+    if tt < 0.0: tt = 0.0
+    if tt > 1.0: tt = 1.0
+    return tt * math.sqrt(1.0 - tt)
+
+@njit(cache=True)
+def dn2(n: int, theta: np.ndarray) -> float:
+    """Distribution function 2 for Distributed model 2."""
+    tt = th(n, theta)
+    if tt < 0.0: tt = 0.0
+    if tt > 1.0: tt = 1.0
+    if tt < 0.5:
+        return math.sqrt(tt)
+    else:
+        return math.sqrt(1.0 - tt)
+
+@njit(cache=True)
+def derivs0(t: float, theta: np.ndarray, C: np.ndarray, gR: np.ndarray) -> np.ndarray:
+    """ Diffusive growth: Model 0 """
+    nMax = len(theta)
+    dThetaDt = np.zeros(nMax)
+    for n in range(nMax):
+        dTheta = th(n-1, theta) - th(n, theta)
+        dThetaDt[n] = (dTheta * gR[n] + 
+                       C[n] * (th(n+1, theta) - th(n+2, theta)) * dTheta - 
+                       C[n] * (th(n, theta) - th(n+1, theta)) * (th(n-2, theta) - th(n-1, theta)))
+    return dThetaDt
+
+@njit(cache=True)
+def derivs1(t: float, theta: np.ndarray, C: np.ndarray, gR: np.ndarray) -> np.ndarray:
+    """ Distributed growth model 1: Model 1 """
+    nMax = len(theta)
+    dThetaDt = np.zeros(nMax)
+    for n in range(nMax):
+        dTheta = th(n-1, theta) - th(n, theta)
+        
+        dn1_n_m1 = dn1(n-1, theta)
+        dn1_n = dn1(n, theta)
+        dn1_n_p1 = dn1(n+1, theta)
+        
+        C_prev = 0.0 if n == 0 else C[n-1]
+        
+        if dn1_n_m1 != 0.0:
+            dTheta -= C_prev * dTheta * dn1_n_m1 / (dn1_n_m1 + dn1_n)
+            
+        if dn1_n > 0.0:
+            dTheta += C[n] * (th(n, theta) - th(n+1, theta)) * dn1_n / (dn1_n + dn1_n_p1)
+            
+        dThetaDt[n] = dTheta * gR[n]
+    return dThetaDt
+
+@njit(cache=True)
+def derivs2(t: float, theta: np.ndarray, C: np.ndarray, gR: np.ndarray) -> np.ndarray:
+    """ Distributed growth model 2: Model 2 """
+    nMax = len(theta)
+    dThetaDt = np.zeros(nMax)
+    for n in range(nMax):
+        dTheta = th(n-1, theta) - th(n, theta)
+        
+        dn2_n_m1 = dn2(n-1, theta)
+        dn2_n = dn2(n, theta)
+        dn2_n_p1 = dn2(n+1, theta)
+        
+        C_prev = 0.0 if n == 0 else C[n-1]
+        
+        if dn2_n_m1 != 0.0:
+            dTheta -= C_prev * dTheta * dn2_n_m1 / (dn2_n_m1 + dn2_n)
+            
+        if dn2_n > 0.0:
+            dTheta += C[n] * (th(n, theta) - th(n+1, theta)) * dn2_n / (dn2_n + dn2_n_p1)
+            
+        dThetaDt[n] = dTheta * gR[n]
+    return dThetaDt
+
+@njit(cache=True)
+def compute_intensity(coverage_history: np.ndarray) -> np.ndarray:
+    """
+    Computes Kinematical Diffracted Intensity.
+    
+    Args:
+        coverage_history (np.ndarray): History matrix of shape (num_time_steps, numLayers).
+        
+    Returns:
+        np.ndarray: Evaluated sequence of intensity values.
+    """
+    num_t = coverage_history.shape[0]
+    num_layers = coverage_history.shape[1]
+    intensity = np.zeros(num_t)
+    
+    PI = math.pi
+    for t in range(num_t):
+        cDI = 1.0 - coverage_history[t, 0]
+        for n in range(1, num_layers):
+            cDI += (coverage_history[t, n-1] - coverage_history[t, n]) * math.cos(n * PI)
+        intensity[t] = cDI * cDI
+    return intensity
+
+@njit(cache=True)
+def compute_rms_roughness(growthTime: np.ndarray, coverage_history: np.ndarray, gR: np.ndarray) -> np.ndarray:
+    """
+    Computes the RMS Roughness over time.
+    
+    Args:
+        growthTime (np.ndarray): Vector of time steps elapsed.
+        coverage_history (np.ndarray): Shape (num_time_steps, numLayers).
+        gR (np.ndarray): Growth rates of the layers.
+        
+    Returns:
+        np.ndarray: Vector of evaluated RMS roughness configurations.
+    """
+    num_t = coverage_history.shape[0]
+    num_layers = coverage_history.shape[1]
+    rms = np.zeros(num_t)
+    
+    for t in range(num_t):
+        sD = (growthTime[t] * gR[0])**2 * (1.0 - coverage_history[t, 0])
+        for n in range(1, num_layers):
+            sD += (n - growthTime[t] * gR[n-1])**2 * (coverage_history[t, n-1] - coverage_history[t, n])
+        rms[t] = math.sqrt(sD)
+    return rms