paddle 1.0.1__tar.gz → 1.1.0__tar.gz

paddle-1.1.0/.gitignore

@@ -0,0 +1,14 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.build/
+dist/
+build/
+.coverage
+htmlcov/
+.pytest_cache/
+.venv/
+.env
+*.log
+

paddle-1.1.0/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: paddle
+Version: 1.1.0
+Summary: Canoe's utility subroutines
+Project-URL: Homepage, https://github.com/chengcli/paddle
+Project-URL: Repository, https://github.com/chengcli/paddle
+Project-URL: Issues, https://github.com/chengcli/paddle/issues
+Author-email: Cheng Li <chengcli@umich.edu>
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Requires-Dist: kintera>=1.0.1
+Requires-Dist: snapy>=0.7.0
+Requires-Dist: torch<=2.7.1,>=2.7.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# paddle
+
+A minimal, utility subroutines for canoe
+
+## Install
+
+```bash
+pip install paddle

paddle-1.1.0/README.md

@@ -0,0 +1,8 @@
+# paddle
+
+A minimal, utility subroutines for canoe
+
+## Install
+
+```bash
+pip install paddle

paddle-1.1.0/pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["hatchling>=1.25"]
+build-backend = "hatchling.build"
+#build-backend = "setuptools.build_meta"
+
+[project]
+name = "paddle"
+version = "1.1.0"
+description = "Canoe's utility subroutines"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [{ name = "Cheng Li", email = "chengcli@umich.edu" }]
+keywords = []
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Topic :: Scientific/Engineering :: Atmospheric Science",
+  "Topic :: Scientific/Engineering :: Physics",
+  "Topic :: Scientific/Engineering :: Astronomy",
+]
+
+dependencies = [
+  "torch>=2.7.0,<=2.7.1",
+  "kintera>=1.0.1",
+  "snapy>=0.7.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=7",
+]
+
+[project.urls]
+Homepage = "https://github.com/chengcli/paddle"
+Repository = "https://github.com/chengcli/paddle"
+Issues = "https://github.com/chengcli/paddle/issues"
+
+[project.scripts]
+paddle = "paddle.__main__:main"

paddle-1.1.0/src/paddle/__init__.py

@@ -0,0 +1,6 @@
+from .setup_profile import setup_profile
+from .write_profile import write_profile
+from .find_init_params import find_init_params
+
+__all__ = ["setup_profile", "write_profile", "find_init_params"]
+__version__ = "1.1.0"

paddle-1.1.0/src/paddle/crm.py

@@ -0,0 +1,76 @@
+import torch
+import math
+import time
+import kintera
+import numpy as np
+from snapy import (
+        index,
+        MeshBlockOptions,
+        MeshBlock,
+        OutputOptions,
+        NetcdfOutput,
+        )
+from kintera import (
+        ThermoOptions,
+        ThermoX,
+        KineticsOptions,
+        Kinetics,
+        )
+
+if __name__ == '__main__':
+    # input file
+    infile = "earth.yaml"
+    device = "cpu"
+
+    # create meshblock
+    op_block = MeshBlockOptions.from_yaml(infile)
+    block = MeshBlock(op_block)
+    block.to(torch.device(device))
+
+    # create thermo module
+    op_thermo = ThermoOptions.from_yaml(infile)
+    thermo_x = ThermoX(op_thermo)
+    thermo_x.to(torch.device(device))
+
+    # create kinetics module
+    op_kinet = KineticsOptions.from_yaml(infile)
+    kinet = Kinetics(op_kinet)
+    kinet.to(torch.device(device))
+
+    # create output fields
+    op_out = OutputOptions().file_basename("earth")
+    out2 = NetcdfOutput(op_out.fid(2).variable("prim"))
+    out3 = NetcdfOutput(op_out.fid(3).variable("uov"))
+    out4 = NetcdfOutput(op_out.fid(4).variable("diag"))
+    outs = [out2, out4]
+
+    # set up initial condition
+    w = setup_initial_condition(block, thermo_x)
+    print("w = ", w[:,0,0,:])
+
+    # integration
+    current_time = 0.0
+    count = 0
+    start_time = time.time()
+    interior = block.part((0, 0, 0))
+    while not block.intg.stop(count, current_time):
+        dt = block.max_time_step()
+        u = block.buffer("hydro.eos.U")
+
+        if count % 1 == 0:
+            print(f"count = {count}, dt = {dt}, time = {current_time}")
+            print("mass = ", u[interior][index.idn].sum())
+
+            for out in outs:
+                out.increment_file_number()
+                out.write_output_file(block, current_time)
+                out.combine_blocks()
+
+        for stage in range(len(block.intg.stages)):
+            block.forward(dt, stage)
+
+        # evolve kinetics
+        u[index.icy:] += evolve_kinetics(block, kinet, thermo_x)
+
+        current_time += dt
+        count += 1

paddle-1.1.0/src/paddle/data/saturn1d.yaml

@@ -0,0 +1,88 @@
+# Saturn Reference Atmosphere Model
+# 
+# Solar abundances relative to H2, enrichments
+#
+# X_He  = 0.195, 0.6955
+# X_CH4 = 5.50e-04, 9.4
+# X_H2O = 1.026e-03, 10.0
+# X_NH3 = 1.352e-04, 3.0
+# X_H2S = 3.10e-05,  3.0
+#
+# Converted to mole fractions
+#
+# X_H2O = 8.91e-03
+# X_NH3 = 3.52e-04
+# X_H2S = 8.08e-05
+#
+# Dry air composition
+#
+# 0.86838 * H2 + 0.11778 * He + 4.49e-03 * CH4
+# {H: 1.755, He: 0.118, C: 4.49e-03}
+
+reference-state:
+  Tref: 0.
+  Pref: 1.e5
+
+species:
+  - name: dry
+    composition: {H: 1.755, He: 0.118, C: 4.49e-03}
+    cv_R: 2.5
+
+  - name: H2O
+    composition: {H: 2, O: 1}
+    cv_R: 2.5
+    u0_R: 0.
+
+  - name: NH3
+    composition: {H: 2, O: 1}
+    cv_R: 2.5
+    u0_R: 0.
+
+  - name: H2S
+    composition: {H: 2, S: 1}
+    cv_R: 2.5
+    u0_R: 0.
+
+  - name: H2O(l)
+    composition: {H: 2, O: 1}
+    cv_R: 9.0
+    u0_R: -3430.
+
+  - name: NH3(s)
+    composition: {H: 2, O: 1}
+    cv_R: 9.0
+    u0_R: -5520.
+
+  - name: NH4SH(s)
+    composition: {N: 1, H: 5, S: 1}
+    cv_R: 9.0
+    u0_R: -1.2e4
+
+geometry:
+  type: cartesian
+  bounds: {x1min: 0., x1max: 600.e3, x2min: -0.5, x2max: 0.5, x3min: -0.5, x3max: 0.5}
+  cells: {nx1: 200, nx2: 1, nx3: 1, nghost: 0}
+
+dynamics:
+  equation-of-state:
+    type: moist-mixture
+    max-iter: 20
+    ftol: 1.e-6
+
+reactions:
+  - equation: H2O => H2O(l)
+    type: nucleation
+    rate-constant: {formula: h2o_ideal}
+
+  - equation: NH3 => NH3(s)
+    type: nucleation
+    rate-constant: {formula: nh3_ideal}
+
+  - equation: NH3 + H2S <=> NH4SH(s)
+    type: nucleation
+    rate-constant: {formula: nh3_h2s_lewis}
+
+boundary-condition:
+  external:
+    x1-inner: reflecting
+    x1-outer: reflecting

paddle-1.1.0/src/paddle/evolve_kinetics.py

@@ -0,0 +1,40 @@
+import torch
+import snapy
+import kintera
+
+def evolve_kinetics(
+    hydro_w: torch.Tensor,
+    block: snapy.MeshBlock, 
+    kinet: kintera.Kinetics, 
+    thermo_x: kintera.ThermoX) -> torch.Tensor:
+    """
+    Evolve the chemical kinetics for one time step using implicit method.
+
+    Args:
+        hydro_w (torch.Tensor): The primitive variables tensor.
+        block (snapy.MeshBlock): The mesh block containing the simulation data.
+        kinet (kintera.Kinetics): The kinetics module for chemical reactions.
+        thermo_x (kintera.ThermoX): The thermodynamics module for computing properties.
+
+    Returns:
+        torch.Tensor: The change in mass density due to chemical reactions.
+    """
+    eos = block.module("hydro.eos")
+    thermo_y = eos.named_modules()["thermo"]
+
+    temp = eos.compute("W->T", (w,))
+    pres = w[index.ipr]
+    xfrac = thermo_y.compute("Y->X", (w[ICY:],))
+    conc = thermo_x.compute("TPX->V", (temp, pres, xfrac))
+    cp_vol = thermo_x.compute("TV->cp", (temp, conc))
+
+    conc_kinet = kinet.options.narrow_copy(conc, thermo_y.options)
+    rate, rc_ddC, rc_ddT = kinet.forward(temp, pres, conc_kinet)
+    jac = kinet.jacobian(temp, conc_kinet, cp_vol, rate, rc_ddC, rc_ddT)
+
+    stoich = kinet.buffer("stoich")
+    del_conc = kintera.evolve_implicit(rate, stoich, jac, dt)
+
+    inv_mu = thermo_y.buffer("inv_mu")
+    del_rho = del_conc / inv_mu[1:].view((1, 1, 1, -1))
+    return del_rho.permute((3, 0, 1, 2))

paddle-1.1.0/src/paddle/find_init_params.py

@@ -0,0 +1,72 @@
+from scipy.interpolate import interp1d
+import snapy
+import numpy as np
+
+from .setup_profile import setup_profile
+
+def find_init_params(
+        block: snapy.MeshBlock,
+        param: dict[str, float],
+        *,
+        target_T: float=300.,
+        target_P: float=1.e5, 
+        method: str="moist-adiabat",
+        max_iter: int=50,
+        ftol: float=1.e-2,
+        verbose: bool=True):
+    """Find initial parameters that yield desired T and P
+
+    Args:
+        block (snapy.MeshBlock): The mesh block to set up.
+        param (dict[str, float]): Initial guess parameters for the adiabat setup.
+            Required keys: Ts, Ps, x<species>, grav.
+        target_T (float, optional): Target temperature in Kelvin. Defaults to 300 K.
+        target_P (float, optional): Target pressure in Pascals. Defaults to 1e5 Pa.
+        method (str, optional): Method for the adiabat setup.
+        max_iter (int, optional): Maximum number of iterations. Defaults to 50.
+        ftol (float, optional): Tolerance for temperature convergence. Defaults to 1e-2 K.
+        verbose (bool, optional): If True, print iteration details. Defaults to True.
+
+    Returns:
+        dict: Dictionary containing the found parameters: Ts, Ps, xH2O, xNH3, xH2S.
+    """
+    count = 0
+    eos = block.hydro.get_eos()
+
+    while count < max_iter:
+        if verbose:
+            print(f"Iteration {count+1}: Trying Ts={param['Ts']}\n")
+
+        # setup profile
+        w = setup_profile(block, param, method=method)
+
+        # calculate temperature
+        temp = eos.compute("W->T", (w,)).squeeze()
+
+        # calculate 1D pressure
+        pres = w[snapy.index.ipr,...].squeeze()
+
+        # temperature function
+        t_func = interp1d(
+            pres.log().cpu().numpy(),
+            temp.log().cpu().numpy(),
+            kind="linear",
+            fill_value="extrapolate")
+
+        temp1 = np.exp(t_func(np.log(target_P)))
+        if verbose:
+            print(f"  At P={target_P:.3e} Pa, T={temp1:.3f} K (target {target_T} K)\n")
+        if abs(temp1 - target_T) < ftol:
+            if verbose:
+                print("Converged! Found parameters:")
+                for key, val in param.items():
+                    print(f"  {key} = {val}")
+                print(f"Matching T = {target_T} K at P = {target_P} Pa")
+            return param
+
+        # adjust Ts using a damped scaling
+        param["Ts"] += (target_T - temp1) * 0.9
+        count += 1
+
+    raise RuntimeError("Failed to converge within the maximum number of iterations.")
+

paddle-1.1.0/src/paddle/setup_profile.py

@@ -0,0 +1,283 @@
+from typing import Tuple
+
+import torch
+import snapy
+import kintera
+
+def integrate_neutral(
+        thermo_x: kintera.ThermoX,
+        temp: torch.Tensor,
+        pres: torch.Tensor,
+        xfrac: torch.Tensor,
+        grav: float,
+        dz: float,
+        max_iter: int = 100
+        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    A neutral density profile assumes no cloud and:
+        
+        (1) dP/dz = -rho*g
+        (2) d(rho)/dz = ...
+
+    In discretized form:
+
+        rho_bar = 0.5 * (rho_old + rho_ad)
+        P_new = P_old - rho_bar * g * dz
+    """
+    conc = thermo_x.compute("TPX->V", [temp, pres, xfrac])
+    rho = thermo_x.compute("V->D", [conc])
+
+    # make an adiabatic step first
+    temp_ad = temp.clone()
+    pres_ad = pres.clone()
+    xfrac_ad = xfrac.clone()
+
+    thermo_x.extrapolate_ad(temp_ad, pres_ad, xfrac_ad, grav, dz)
+    conc_ad = thermo_x.compute("TPX->V", [temp_ad, pres_ad, xfrac_ad])
+    rho_ad = thermo_x.compute("V->D", [conc_ad])
+    rho_bar = 0.5 * (rho + rho_ad)
+
+    pres2 = pres - rho_bar * grav * dz
+
+    # initial guess
+    temp2 = temp_ad.clone()
+    count = 0
+    while count < max_iter:
+        xfrac2 = xfrac_ad.clone()
+
+        # equilibrate clouds
+        thermo_x.forward(temp2, pres2, xfrac2)
+
+        # drop clouds fractions
+        for cid in thermo_x.options.cloud_ids():
+            xfrac2[..., cid] = 0.0
+        # renormalize mole fractions
+        xfrac2 /= xfrac.sum(dim=-1, keepdim=True)
+
+        conc2 = thermo_x.compute("TPX->V", [temp2, pres2, xfrac2])
+        rho2 = thermo_x.compute("V->D", [conc2])
+
+        if torch.allclose(rho2, rho_ad):
+            break
+
+        temp2 -= temp2 * (rho_ad - rho2) / rho2
+        count += 1
+
+    if count == max_iter:
+        raise RuntimeError("neutral density integration did not converge.")
+
+    return temp2, pres2, xfrac2
+
+def integrate_dry_adiabat(
+        thermo_x: kintera.ThermoX,
+        temp: torch.Tensor,
+        pres: torch.Tensor,
+        xfrac: torch.Tensor,
+        grav: float,
+        dz: float,
+        max_iter: int = 100
+        ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    A dry adiabatic profile assumes no cloud and:
+
+        (1) dT/dz = -g/cp
+        (2) dP/dz = -rho*g
+
+    In discretized form:
+    
+        cp_bar = 0.5 * (cp(T_old) + cp(T_new))
+        T_new = T_old - g/bar * dz
+        rho_bar = 0.5 * (rho_old + rho_new)
+        P_new = P_old - rho_bar * g * dz
+        
+    """
+    conc1 = thermo_x.compute("TPX->V", [temp, pres, xfrac])
+    cp1 = thermo_x.compute("TV->cp", [temp, conc1]) / conc1.sum(-1)
+    rho1 = thermo_x.compute("V->D", [conc1])
+    mmw1 = (thermo_x.mu * xfrac).sum(-1)
+
+    # initial guess
+    temp2 = temp - grav * mmw1 * dz / cp1
+    pres2 = pres - rho1 * grav * dz
+
+    count = 0
+    while count < max_iter:
+        xfrac2 = xfrac.clone()
+
+        # equilibrate clouds
+        thermo_x.forward(temp2, pres2, xfrac2)
+
+        # drop clouds fractions
+        for cid in thermo_x.options.cloud_ids():
+            xfrac2[..., cid] = 0.0
+        # renormalize mole fractions
+        xfrac2 /= xfrac.sum(dim=-1, keepdim=True)
+
+        conc2 = thermo_x.compute("TPX->V", [temp2, pres2, xfrac2])
+        cp2 = thermo_x.compute("TV->cp", [temp2, conc2]) / conc2.sum(-1)
+        rho2 = thermo_x.compute("V->D", [conc2])
+        mmw2 = (thermo_x.mu * xfrac2).sum(-1)
+
+        cp_bar = 0.5 * (cp1 / mmw1 + cp2 / mmw2)
+        rho_bar = 0.5 * (rho1 + rho2)
+
+        temp_new = temp - grav * dz / cp_bar
+        pres_new = pres - rho_bar * grav * dz
+
+        if torch.allclose(temp_new, temp2) and torch.allclose(pres_new, pres2):
+            break
+
+        temp2 = temp_new
+        pres2 = pres_new
+        count += 1
+
+    if count == max_iter:
+        raise RuntimeError("Dry adiabat integration did not converge.")
+
+    return temp2, pres2, xfrac2
+
+def setup_profile(
+        block: snapy.MeshBlock,
+        param: dict[str, float] = {},
+        method: str = "moist-adiabat"
+        ) -> torch.Tensor:
+    """
+    Set up an adiabatic initial condition for the mesh block.
+
+    This function initializes the primitive variables in the mesh block
+    and returns the initialized tensor.
+
+    Args:
+        block (snapy.MeshBlock): The mesh block to set up.
+        param (dict[str, float], optional): Parameters for the adiabat setup. Defaults to {}.
+        method (str, optional): Method for the adiabat setup. Choose between
+            (1) "dry-adiabat"
+            (2) "moist-adiabat"
+            (3) "isothermal"
+            (4) "pseudo-adiabat"
+            (5) "neutral"
+            Defaults to "moist-adiabat".
+
+        Required parameters in `param`:
+            Ts (float): Surface temperature in Kelvin. Default is 300 K.
+            Ps (float): Surface pressure in Pascals. Default is 1e5 Pa.
+            x<species> (float): Mole fraction of a specific species (e.g., xH2O for
+            water vapor). Default is 0.0.
+            grav (float): Gravitational acceleration in m/s^2. Default is 9.8 m/s^2.
+
+    Returns:
+        torch.Tensor: The initialized primitive variables tensor.
+    """
+
+    # check method
+    valid_methods = [
+        "dry-adiabat",
+        "moist-adiabat",
+        "isothermal",
+        "pseudo-adiabat",
+        "neutral"
+    ]
+
+    if method not in valid_methods:
+        raise ValueError(f"Invalid method '{method}'. Choose from {valid_methods}.")
+
+    Ts = param.get("Ts", 300.)
+    Ps = param.get("Ps", 1.e5)
+    grav = param.get("grav", 9.8)
+    Tmin = param.get("Tmin", 0.)
+
+    # get handles to modules
+    coord = block.module("hydro.coord")
+    thermo_y = block.module("hydro.eos.thermo")
+
+    # get coordinates
+    x3v, x2v, x1v = torch.meshgrid(
+        coord.buffer("x3v"), coord.buffer("x2v"), coord.buffer("x1v"), indexing="ij"
+    )
+
+    # handling mole fractions
+    thermo_x = kintera.ThermoX(thermo_y.options)
+    thermo_x.to(dtype=x1v.dtype, device=x1v.device)
+
+    # get dimensions
+    nc3, nc2, nc1 = x1v.shape
+    ny = len(thermo_y.options.species()) - 1
+    nvar = 5 + ny
+
+    w = torch.zeros((nvar, nc3, nc2, nc1),
+                    dtype=x1v.dtype, device=x1v.device)
+
+    temp = Ts * torch.ones((nc3, nc2), dtype=w.dtype, device=w.device)
+    pres = Ps * torch.ones((nc3, nc2), dtype=w.dtype, device=w.device)
+    xfrac = torch.zeros((nc3, nc2, ny + 1), dtype=w.dtype, device=w.device)
+
+    for name in thermo_y.options.species():
+        index = thermo_y.options.species().index(name)
+        xfrac[..., index] = param.get(f"x{name}", 0.0)
+
+    # dry air mole fraction
+    xfrac[..., 0] = 1. - xfrac[..., 1:].sum(dim=-1)
+
+    # start and end indices for the vertical direction
+    # excluding ghost cells
+    ifirst = coord.ifirst()
+    ilast = coord.ilast()
+
+    # vertical grid distance of the first cell
+    dz = coord.buffer("dx1f")[ifirst]
+
+    # half a grid to cell center
+    thermo_x.extrapolate_ad(temp, pres, xfrac, grav, dz / 2.);
+
+    # adiabatic extrapolation
+    if method == "isothermal":
+        i_isothermal = ifirst
+        ifirst = ilast
+    else:
+        i_isothermal = ilast
+
+    for i in range(ifirst, ilast):
+        # drop clouds fractions
+        if method.split("-")[0] != "moist":
+            for cid in thermo_x.options.cloud_ids():
+                xfrac[..., cid] = 0.0
+            # renormalize mole fractions
+            xfrac /= xfrac.sum(dim=-1, keepdim=True)
+        conc = thermo_x.compute("TPX->V", [temp, pres, xfrac])
+
+        w[snapy.index.ipr, ..., i] = pres;
+        w[snapy.index.idn, ..., i] = thermo_x.compute("V->D", [conc])
+        w[snapy.index.icy:, ...,i] = thermo_x.compute("X->Y", [xfrac])
+
+        dz = coord.buffer("dx1f")[i]
+        if method.split("-")[0] == "dry":
+            temp, pres, xfrac = integrate_dry_adiabat(thermo_x, temp, pres, xfrac, grav, dz);
+        elif method.split("-")[0] == "neutral":
+            temp, pres, xfrac = integrate_neutral(thermo_x, temp, pres, xfrac, grav, dz);
+        else:
+            thermo_x.extrapolate_ad(temp, pres, xfrac, grav, dz);
+
+        if torch.any(temp < Tmin):
+            i_isothermal = i + 1
+            break
+
+    # isothermal extrapolation
+    for i in range(i_isothermal, ilast):
+        # drop clouds fractions
+        if method.split("-")[0] != "moist":
+            for cid in thermo_x.options.cloud_ids():
+                xfrac[..., cid] = 0.0
+            # renormalize mole fractions
+            xfrac /= xfrac.sum(dim=-1, keepdim=True)
+
+        mu = (thermo_x.mu * xfrac).sum(-1)
+        dz = coord.buffer("dx1f")[i]
+        pres *= torch.exp(-grav * mu * dz / (kintera.constants.Rgas * temp))
+        conc = thermo_x.compute("TPX->V", [temp, pres, xfrac])
+        w[snapy.index.ipr, ..., i] = pres
+        w[snapy.index.idn, ..., i] = thermo_x.compute("V->D", [conc])
+        w[snapy.index.icy:, ..., i] = thermo_x.compute("X->Y", [xfrac])
+
+    # initialize hydro state
+    block.initialize({"hydro_w": w})
+    return w