py-growth-RHEED 1.0.0__tar.gz

py_growth_rheed-1.0.0/LICENSE

@@ -0,0 +1,11 @@
+Copyright (c) 2026, Andrzej Daniluk & Bartek Daniluk
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

py_growth_rheed-1.0.0/PKG-INFO

@@ -0,0 +1,153 @@
+Metadata-Version: 2.4
+Name: py-growth-RHEED
+Version: 1.0.0
+Summary: A Python package for simulations of RHEED intensity oscillations within the kinematical approximation
+Author: Bartek Daniluk
+License: BSD 3-Clause
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Intended Audience :: Science/Research
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2.5.0,>=2.0.0
+Requires-Dist: numba>=0.60.0
+Requires-Dist: jax>=0.4.30
+Requires-Dist: jaxlib>=0.4.30
+Requires-Dist: plotly>=5.24.0
+Requires-Dist: jupyter
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# PY_GROWTH Simulation Package
+
+<div align="center">
+  <img src="https://img.shields.io/badge/Python-3.10+-blue.svg" alt="Python Version"/>
+  <img src="https://img.shields.io/badge/JAX-Accelerated-orange.svg" alt="JAX Accelerated"/>
+  <img src="https://img.shields.io/badge/Numba-JIT-green.svg" alt="Numba JIT"/>
+  <img src="https://img.shields.io/badge/Status-Production-success.svg" alt="Production Ready"/>
+</div>
+
+**PY_GROWTH** is a modern, high-performance Python package for simulating layer coverages during the growth of thin epitaxial films and the corresponding **RHEED (Reflection High-Energy Electron Diffraction)** intensities within the kinematical approximation. 
+
+This package translates, refactors, and massively accelerates legacy C++ simulation algorithms using state-of-the-art numerical engines (`NumPy`, `Numba LLVM`, and `JAX/XLA`), enabling rapid programmatic experimentation and data plotting.
+
+---
+
+## 1. Quick Start & Tutorial
+
+Because PY_GROWTH is built dynamically for modern Python workflows, installing the package and running a simulation requires fewer than 5 lines of code.
+
+### Installation
+
+The package can be installed globally into any Virtual Environment or Jupyter Server utilizing the newly added `setup.py`:
+```bash
+# Install the package globally via PyPI
+pip install py-RHEED
+```
+This automatically resolves all physics dependencies like `jax`, `numba`, and visualization libraries like `plotly`.
+
+### Interactive Jupyter Workflow
+
+1. Prepare your input constraints inside the `exampleData/inputData.dat` file.
+2. Open Jupyter: `python -m notebook`
+3. Enter the unified execution pipeline!
+
+```python
+import plotly.graph_objects as go
+from PY_GROWTH.numpy import run_simulation
+from PY_GROWTH.io import load_input_data
+import os
+
+# 1. Load mathematical constraints safely from your file
+file_path = os.path.join("exampleData", "inputData.dat")
+model_type, num_layers, t_max, num_intervals, linear_rel_an_kn, linear_rel_grn, an_kn, growth_rates = load_input_data(file_path)
+
+# 2. Execute the Adaptive Runge-Kutta solver
+results = run_simulation(model_type, num_layers, t_max, num_intervals, an_kn, growth_rates)
+
+# 3. Visualize RHEED Oscillations immediately
+fig = go.Figure()
+fig.add_trace(go.Scatter(x=results['time'], y=results['intensity'], mode='lines', name='Intensity'))
+fig.update_layout(title="RHEED Intensity Oscillations", xaxis_title="Time (s)", yaxis_title="RHEED Intensity")
+fig.show()
+
+# 4. Optional: Save raw data identically to C++ flushing
+import pandas as pd
+import numpy as np
+df = pd.DataFrame(np.column_stack([results['time'], results['intensity'], results['y']]))
+df.to_csv("RHEED_Oscillations_Export.csv", index=False)
+```
+
+See the `examples/` directory for full script files using `JAX` and `Numba` compilers!
+
+---
+
+## 2. Theory & Background
+
+PY_GROWTH provides a deterministic, rigorous numerical solution for initial value problems analyzing nonlinear differential equations. The user provides constraints dictating how growth parameters ($A_n$ or $k_n$) act relative to linear intervals.
+
+The software computes three primary physical models:
+* **Model 0 (Diffusive Growth):** Simulates standard continuous step-flow and localized layer progression.
+* **Model 1 (Distributed Growth Type 1):** Introduces probabilistic dispersion mechanics across surface coverages.
+* **Model 2 (Distributed Growth Type 2):** Alters distribution bounds to simulate chaotic island generation.
+
+### The Input Data Schema
+
+The `inputData.dat` file adheres to a strict C++-legacy schema. It expects floating-point matrices defining the system exactly in this sequence:
+1. `modelType`: 0, 1, or 2 limits.
+2. `numLayers`: The number of discrete vertical atom sequences evaluated simultaneously.
+3. `tMax`: Upper limit of the growth time interval.
+4. `num_intervals`: Granularity of output data.
+5. Boolean flags dictating continuous versus fixed constraints.
+6. The $A_n$ (or $k_n$) arrays mapped individually to each layer boundary limits.
+7. Growth rates representing $1/\tau_n \text{ MLs}^{-1}$.
+
+--- 
+
+## 3. Academic Citations
+
+If incorporating **PY_GROWTH** into primary research publications or academic frameworks, please consider citing the canonical compiler topologies central to our physics translation engines:
+
+* **The ODE Solver (Cash-Karp Runge-Kutta 5th-Order):** 
+  Cash, J. R., & Karp, A. H. (1990). *A variable order Runge-Kutta method for initial value problems with rapidly varying right-hand sides.* ACM Transactions on Mathematical Software (TOMS), 16(3), 201-222.
+* **The Numba LLVM Hardware Architecture:** 
+  Lam, S. K., Pitrou, A., & Seibert, S. (2015). *Numba: A LLVM-based Python JIT compiler.* In Proceedings of the Second Workshop on the LLVM Compiler Infrastructure in HPC (pp. 1-6).
+* **The Google JAX/XLA Accelerator Integration:** 
+  Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., & Zhang, Q. (2018). *JAX: composable transformations of Python+NumPy programs.* http://github.com/google/jax
+---
+
+## 4. PyPI Publication (Developer Notes)
+
+Because PY_GROWTH utilizes a standard setup.py build configuration, it is fully capable of being hosted on the global **Python Package Index (PyPI)**. 
+
+To publish updates to the global repository (allowing users worldwide to use pip install py-growth), authenticate your PyPI account and execute:
+`ash
+# 1. Install build tools
+pip install build twine
+
+# 2. Compile the package into a wheel distribution
+python -m build
+
+# 3. Securely upload to PyPI servers
+twine upload dist/*
+`
+
+---
+
+## License
+This project is licensed under the **BSD 3-Clause License** - see the [LICENSE](LICENSE) file for details.
+
+
+## 4. License
+
+This project is licensed under the **BSD 3-Clause License** - see the LICENSE file for details. 
+Free for academic, research, and commercial use.

py_growth_rheed-1.0.0/PY_GROWTH/__init__.py

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

py_growth_rheed-1.0.0/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_rheed-1.0.0/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_rheed-1.0.0/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_rheed-1.0.0/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_rheed-1.0.0/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_rheed-1.0.0/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
+    }