jaxprop 0.3.1__tar.gz

jaxprop-0.3.1/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Roberto Agromayor
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jaxprop-0.3.1/PKG-INFO

@@ -0,0 +1,49 @@
+Metadata-Version: 2.1
+Name: jaxprop
+Version: 0.3.1
+Summary: JAX-compatible thermodynamic property calculations.
+Home-page: https://github.com/turbo-sim/jaxprop
+License: MIT
+Author: Roberto Agromayor
+Requires-Python: >=3.11,<3.14
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: CoolProp (>=7.0.0,<8.0.0)
+Requires-Dist: diffrax (>=0.7.0,<0.8.0)
+Requires-Dist: equinox (>=0.12.0,<0.13.0)
+Requires-Dist: jax (>=0.6.0,<0.7.0)
+Requires-Dist: jaxlib (>=0.6.0,<0.7.0)
+Requires-Dist: matplotlib (>=3.10.6,<4.0.0)
+Requires-Dist: numpy (>=2.3.0,<3.0.0)
+Requires-Dist: optimistix (>=0.0.10,<0.0.11)
+Requires-Dist: pysolver_view (>=0.6.8,<0.7.0)
+Requires-Dist: scipy (>=1.16.1,<2.0.0)
+Project-URL: Documentation, https://github.com/turbo-sim/jaxpconrop
+Project-URL: Repository, https://github.com/turbo-sim/jaxprop
+Description-Content-Type: text/markdown
+
+# JAXprop
+
+`jaxprop` provides JAX-compatible thermodynamic property calculations with support for automatic differentiation, vectorization, and JIT compilation.  
+
+🔗 **Docs**: [turbo-sim.github.io/jaxprop](https://turbo-sim.github.io/jaxprop/)  
+📦 **PyPI**: [pypi.org/project/jaxprop](https://pypi.org/project/jaxprop/)
+
+**Note**: This project is based on the [CoolProp](https://www.coolprop.org) library but is not affiliated with or endorsed by the CoolProp project.
+
+## Key features
+
+- Compute and plot phase envelopes and spinodal lines for pure fluids.
+- Evaluate thermodynamic properties from Helmholtz energy–based equations of state, including metastable states inside the two-phase region.
+- Perform flash calculations for any input pair with a custom solver and user-defined initial guesses.
+- Work with structured property dictionaries and immutable `FluidState` objects.
+- Evaluate properties over arrays of input conditions for efficient parametric studies and plotting.
+- Full JAX compatibility: supports `jit`, `grad`, `vmap`, and parallel evaluation.
+
+## Installation
+
+```bash
+pip install jaxprop
+

jaxprop-0.3.1/README.md

@@ -0,0 +1,22 @@
+# JAXprop
+
+`jaxprop` provides JAX-compatible thermodynamic property calculations with support for automatic differentiation, vectorization, and JIT compilation.  
+
+🔗 **Docs**: [turbo-sim.github.io/jaxprop](https://turbo-sim.github.io/jaxprop/)  
+📦 **PyPI**: [pypi.org/project/jaxprop](https://pypi.org/project/jaxprop/)
+
+**Note**: This project is based on the [CoolProp](https://www.coolprop.org) library but is not affiliated with or endorsed by the CoolProp project.
+
+## Key features
+
+- Compute and plot phase envelopes and spinodal lines for pure fluids.
+- Evaluate thermodynamic properties from Helmholtz energy–based equations of state, including metastable states inside the two-phase region.
+- Perform flash calculations for any input pair with a custom solver and user-defined initial guesses.
+- Work with structured property dictionaries and immutable `FluidState` objects.
+- Evaluate properties over arrays of input conditions for efficient parametric studies and plotting.
+- Full JAX compatibility: supports `jit`, `grad`, `vmap`, and parallel evaluation.
+
+## Installation
+
+```bash
+pip install jaxprop

jaxprop-0.3.1/jaxprop/__init__.py

@@ -0,0 +1,46 @@
+# Highlight exception messages
+# https://stackoverflow.com/questions/25109105/how-to-colorize-the-output-of-python-errors-in-the-gnome-terminal/52797444#52797444
+try:
+    import IPython.core.ultratb
+except ImportError:
+    # No IPython. Use default exception printing.
+    pass
+else:
+    import sys
+    sys.excepthook = IPython.core.ultratb.FormattedTB(color_scheme='linux', call_pdb=False)
+
+
+import os
+os.environ["JAX_PLATFORM_NAME"] = "cpu"
+import jax
+jax.config.update("jax_enable_x64", True)
+
+
+
+from .graphics import *
+from .utils import *
+from .helpers_jax import *
+from .helpers_coolprop import *
+
+# Import subpackages
+from . import coolprop
+from . import perfect_gas
+# from . import bicubic
+
+# Import API classes
+from .perfect_gas import FluidPerfectGas
+from .coolprop import Fluid, FluidJAX
+
+
+from . import components
+
+
+# Package info
+__version__ = "0.3.1"
+PACKAGE_NAME = "jaxprop"
+URL_GITHUB = "https://github.com/turbo-sim/jaxprop"
+URL_DOCS = "https://turbo-sim.github.io/jaxprop/"
+URL_PYPI = "https://pypi.org/project/jaxprop/"
+URL_DTU = "https://thermalpower.dtu.dk/"
+BREAKLINE = 80 * "-"
+

jaxprop-0.3.1/jaxprop/bicubic/__init__.py

@@ -0,0 +1,3 @@
+from .generate_tables import *
+from .bicubic_interpolant_property import *
+from .jax_bicubic_HEOS_interpolation_1 import *

jaxprop-0.3.1/jaxprop/bicubic/bicubic_interpolant_property.py

@@ -0,0 +1,69 @@
+import numpy as np
+# import jax.numpy as jnp
+from ..jax_import import jnp
+import pickle
+
+from .jax_bicubic_HEOS_interpolation_1 import compute_bicubic_coefficients_of_ij, bicubic_interpolant
+
+
+def bicubic_interpolant_property(h, P, table):
+    """
+    Interpolates all properties at a given enthalpy (h) and pressure (P) using bicubic interpolation.
+
+    Args:
+        h (float): Enthalpy in J/kg
+        P (float): Pressure in Pa
+        table_path (str): Path to saved property table (dict-based .pkl)
+
+    Returns:
+        dict: Dictionary of interpolated property values at (h, P)
+    """
+    h_vals = jnp.array(table['h'])
+    P_vals = jnp.array(table['P'])
+
+    Nh, Np = len(h_vals), len(P_vals)
+    hmin, hmax = float(h_vals[0]), float(h_vals[-1])
+    Lmin, Lmax = float(jnp.log(P_vals[0])), float(jnp.log(P_vals[-1]))
+
+    logP = jnp.log(P)
+
+    # Identify cell (i, j)
+    i = int((h - hmin) / (hmax - hmin) * (Nh - 1))
+    j = int((logP - Lmin) / (Lmax - Lmin) * (Np - 1))
+
+    i = np.clip(i, 0, Nh - 2)
+    j = np.clip(j, 0, Np - 2)
+
+    deltah = float(h_vals[1] - h_vals[0])
+    deltaL = float(jnp.log(P_vals[1]) - jnp.log(P_vals[0]))
+
+    interpolated_props = {}
+
+    for prop, prop_data in table.items():
+        if prop in ['h', 'P']:
+            continue  # skip grid axes
+
+        f_grid = prop_data['value']
+        fx_grid = prop_data['d_dh']
+        fy_grid = prop_data['d_dP']
+        fxy_grid = prop_data['d2_dhdP']
+        
+        coeffs_local = compute_bicubic_coefficients_of_ij(
+            i, j,
+            f_grid,
+            fx_grid * deltah,
+            fy_grid * deltaL,
+            fxy_grid * deltah * deltaL
+        )
+
+        coeffs = jnp.zeros((Nh, Np, 16), dtype=jnp.float64)
+        coeffs = coeffs.at[i, j, :].set(coeffs_local)
+
+        val = bicubic_interpolant(
+            h, P, h_vals, jnp.log(P_vals), coeffs,
+            Nh, Np, hmin, hmax, Lmin, Lmax
+        )
+
+        interpolated_props[prop] = float(val)
+
+    return interpolated_props

jaxprop-0.3.1/jaxprop/bicubic/generate_tables.py

@@ -0,0 +1,85 @@
+import numpy as np
+
+from ..jax_import import jnp
+# import jax.numpy as jnp
+import CoolProp.CoolProp as cp
+# import turboflow as tf
+import os
+import pickle
+from jaxprop.fluid_properties import Fluid
+
+def generate_property_table(hmin, hmax, Pmin, Pmax, fluid_name, Nh, Np, outdir='fluid_tables'):
+    fluid = Fluid(fluid_name)
+    # fluid = tf.Fluid(fluid_name)
+    h_vals = jnp.linspace(hmin, hmax, Nh)
+    Lmin=jnp.log(Pmin) # Log of P
+    Lmax=jnp.log(Pmax) # Log of P
+    P_vals = jnp.linspace(Lmin, Lmax, Np)
+
+    deltah = float(h_vals[1] - h_vals[0])
+    deltaL = P_vals[1]-P_vals[0]
+    eps_h = 0.001 * deltah
+    eps_P = 1e-6 * float(Pmin)
+
+    properties = {
+        'T': 'T',       # Temperature [K]
+        'd': 'D',       # Density [kg/m³]
+        's': 'S',       # Entropy [J/kg/K]
+        'mu': 'V',      # Viscosity [Pa·s]
+        'k': 'L',       # Thermal conductivity [W/m/K]
+    }
+
+    # Initialize property grids
+    table = {
+        'h': np.array(h_vals),
+        'P': np.array(jnp.exp(P_vals))
+    }
+
+    for key in properties:
+        table[key] = {
+            'value': np.zeros((Nh, Np), dtype=np.float64),
+            'd_dh':  np.zeros((Nh, Np), dtype=np.float64),
+            'd_dP':  np.zeros((Nh, Np), dtype=np.float64),
+            'd2_dhdP': np.zeros((Nh, Np), dtype=np.float64),
+        }
+
+    # Loop over grid and populate values
+    for i, h in enumerate(h_vals):
+        for j, P in enumerate(P_vals):
+            hf = float(h)
+            Pf = jnp.exp(P)
+
+            try:        
+                f_0 = fluid.get_state(cp.HmassP_INPUTS, float(hf), float(Pf))
+                f_h = fluid.get_state(cp.HmassP_INPUTS, float(hf + eps_h), float(Pf))
+                f_p = fluid.get_state(cp.HmassP_INPUTS, float(hf), float(Pf + eps_P))
+                f_hp = fluid.get_state(cp.HmassP_INPUTS, float(hf + eps_h), float(Pf + eps_P))
+            # Uncomment below if using turboflow
+                # f_0 = tf.get_props_custom_jvp(fluid, cp.HmassP_INPUTS, hf, Pf)
+                # f_h = tf.get_props_custom_jvp(fluid, cp.HmassP_INPUTS, hf + eps_h, Pf)
+                # f_p = tf.get_props_custom_jvp(fluid, cp.HmassP_INPUTS, hf, Pf + eps_P)
+                # f_hp = tf.get_props_custom_jvp(fluid, cp.HmassP_INPUTS, hf + eps_h, Pf + eps_P)
+            except Exception:
+                continue  # Skip invalid points
+
+            for key in properties:
+                val = f_0[key]
+                dval_dh = (f_h[key] - f_0[key]) / eps_h
+                dval_dP = (f_p[key] - f_0[key]) / eps_P
+                d2val_dhdP = (f_hp[key] - f_h[key] - f_p[key] + f_0[key]) / (eps_h * eps_P)
+
+                table[key]['value'][i, j] = val
+                table[key]['d_dh'][i, j] = dval_dh
+                table[key]['d_dP'][i, j] = dval_dP
+                table[key]['d2_dhdP'][i, j] = d2val_dhdP
+
+    # Save as pickle only (most useful for JAX processing)
+    os.makedirs(outdir, exist_ok=True)
+    pkl_path = os.path.join(outdir, f'{fluid_name}_{Nh}_x_{Np}.pkl')
+
+    with open(pkl_path, 'wb') as f:
+        pickle.dump(table, f)
+
+    print(f" Saved the table to:\n -> Pickle: {pkl_path}")
+
+    return table

jaxprop-0.3.1/jaxprop/bicubic/jax_bicubic_HEOS_interpolation_1.py

@@ -0,0 +1,292 @@
+# ======================== Imports ========================
+# import os
+# import psutil
+import time
+# import jax
+# import jax.numpy as jnp
+# from jax import jit
+
+from ..jax_import import jax, jnp, jit
+# from jax.experimental import mesh_utils
+# from jax.sharding import Mesh, PartitionSpec, NamedSharding
+import CoolProp.CoolProp as cp
+import numpy as np
+import matplotlib.pyplot as plt
+from functools import partial
+
+# ======================== Config ========================
+# NCORES = psutil.cpu_count(logical=False)
+# os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={NCORES}"
+# jax.config.update("jax_enable_x64", True)
+
+# Global precision
+float64 = jnp.dtype("float64")
+complex128 = jnp.dtype("complex128")
+
+# =================== Functions to Export ===================
+# @jax.jit
+def compute_bicubic_coefficients_of_ij(i, j, f, fx, fy, fxy):
+    #xx=f(0,0)&f(1,0)&f(0,1)&f(1,1)&f_x(0,0)&f_x(1,0)&f_x(0,1)&f_x(1,1)&f_y(0,0)&f_y(1,0)&f_y(0,1)&f_y(1,1)&f_{xy}(0,0)&f_{xy}(1,0)&f_{xy}(0,1)&f_{xy}(1,1)
+    xx=[f[i,j],f[i+1,j],f[i,j+1],f[i+1,j+1],fx[i,j],fx[i+1,j],fx[i,j+1],fx[i+1,j+1],fy[i,j],fy[i+1,j],fy[i,j+1],fy[i+1,j+1],fxy[i,j],fxy[i+1,j],fxy[i,j+1],fxy[i+1,j+1]]
+    A=[ [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0 ],
+        [ 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0 ],
+        [ -3., 3., 0., 0., -2., -1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0 ],
+        [ 2., -2., 0., 0., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0 ],
+        [ 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0 ],
+        [ 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0 ],
+        [ 0., 0., 0., 0., 0., 0., 0., 0., -3., 3., 0., 0., -2., -1., 0., 0 ],
+        [ 0., 0., 0., 0., 0., 0., 0., 0., 2., -2., 0., 0., 1., 1., 0., 0 ],
+        [ -3., 0., 3., 0., 0., 0., 0., 0., -2., 0., -1., 0., 0., 0., 0., 0 ],
+        [ 0., 0., 0., 0., -3., 0., 3., 0., 0., 0., 0., 0., -2., 0., -1., 0 ],
+        [ 9., -9., -9., 9., 6., 3., -6., -3., 6., -6., 3., -3., 4., 2., 2., 1 ],
+        [ -6., 6., 6., -6., -3., -3., 3., 3., -4., 4., -2., 2., -2., -2., -1., -1 ],
+        [ 2., 0., -2., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0 ],
+        [ 0., 0., 0., 0., 2., 0., -2., 0., 0., 0., 0., 0., 1., 0., 1., 0 ],
+        [ -6., 6., 6., -6., -4., -2., 4., 2., -3., 3., -3., 3., -2., -1., -2., -1 ],
+        [ 4., -4., -4., 4., 2., 2., -2., -2., 2., -2., 2., -2., 1., 1., 1., 1]]
+    return jnp.matmul(jnp.array(A,dtype=f.dtype),jnp.array(xx,dtype=f.dtype))
+
+
+# def bicubic_interpolant(h_vals, P_vals, coeffs, hmin, hmax, Lmin, Lmax, Nh, Np):
+#     """
+#     Create a bicubic interpolant function using precomputed coefficients.
+#     """
+#     # Ensure that coeffs is in the right shape (Nh, Np, 16)
+#     assert coeffs.shape == (Nh, Np, 16), f"Expected coeffs to have shape (Nh, Np, 16), but got {coeffs.shape}"
+    
+#     # Normalized grid values
+#     def normalize(value, min_val, max_val):
+#         return (value - min_val) / (max_val - min_val)
+
+#     # The actual interpolant function, it takes h and P as arguments
+#     @jit
+#     def interpolant_fn(h, P):
+#         # Ensure h_vals and P_vals are 1D arrays
+#         h_vals_flat = jnp.ravel(h_vals)  # Flatten the h_vals to 1D
+#         P_vals_flat = jnp.ravel(P_vals)  # Flatten the P_vals to 1D
+        
+#         # Normalize h and P
+#         norm_h = normalize(h, hmin, hmax)
+#         norm_P = normalize(P, Lmin, Lmax)
+
+#         # Identify the grid points surrounding h, P
+#         i = jnp.clip(jnp.searchsorted(h_vals_flat, h) - 1, 0, Nh - 2)
+#         j = jnp.clip(jnp.searchsorted(P_vals_flat, P) - 1, 0, Np - 2)
+
+#         # Extract the coefficients for the surrounding grid points
+#         coeff = coeffs[i, j]
+
+#         # Interpolate in both directions
+#         h_diff = norm_h - h_vals_flat[i]
+#         P_diff = norm_P - P_vals_flat[j]
+
+#         # Calculate the interpolant using the bicubic coefficients
+#         result = (
+#             coeff[0] + coeff[1] * h_diff + coeff[2] * P_diff + coeff[3] * h_diff * P_diff +
+#             coeff[4] * h_diff**2 + coeff[5] * h_diff * P_diff**2 +
+#             coeff[6] * P_diff**2 + coeff[7] * h_diff**2 * P_diff + 
+#             coeff[8] * h_diff**3 + coeff[9] * h_diff**2 * P_diff +
+#             coeff[10] * h_diff * P_diff**2 + coeff[11] * h_diff**3 * P_diff + 
+#             coeff[12] * P_diff**3 + coeff[13] * h_diff * P_diff**3 +
+#             coeff[14] * h_diff**3 * P_diff**2 + coeff[15] * h_diff**2 * P_diff**3
+#         )
+
+#         return result
+
+#     return interpolant_fn
+
+# @partial(jit, static_argnums=(5, 6))  # Nh and Np are static # static arguments are all except h, P
+# def bicubic_interpolant(h, P, h_vals, P_vals, coeffs, Nh, Np, hmin, hmax, Lmin, Lmax):
+#     """
+#     Evaluate the bicubic interpolant at (h, P) using precomputed coefficients.
+#     """
+#     # Normalize
+#     norm_h = (h - hmin) / (hmax - hmin)
+#     norm_P = (P - Lmin) / (Lmax - Lmin)
+
+#     # Flatten grid arrays
+#     h_vals_flat = jnp.ravel(h_vals)
+#     P_vals_flat = jnp.ravel(P_vals)
+
+#     # Find surrounding grid indices
+#     i = jnp.clip(jnp.searchsorted(h_vals_flat, h) - 1, 0, Nh - 2)
+#     j = jnp.clip(jnp.searchsorted(P_vals_flat, P) - 1, 0, Np - 2)
+
+#     # Relative differences in normalized space
+#     h_base = (h_vals_flat[i] - hmin) / (hmax - hmin)
+#     P_base = (P_vals_flat[j] - Lmin) / (Lmax - Lmin)
+#     h_diff = norm_h - h_base
+#     P_diff = norm_P - P_base
+
+#     # Fetch bicubic coefficients
+#     coeff = coeffs[i, j]
+
+#     # Evaluate bicubic polynomial
+#     result = (
+#         coeff[0] + coeff[1] * h_diff + coeff[2] * P_diff + coeff[3] * h_diff * P_diff +
+#         coeff[4] * h_diff**2 + coeff[5] * h_diff * P_diff**2 +
+#         coeff[6] * P_diff**2 + coeff[7] * h_diff**2 * P_diff +
+#         coeff[8] * h_diff**3 + coeff[9] * h_diff**2 * P_diff +
+#         coeff[10] * h_diff * P_diff**2 + coeff[11] * h_diff**3 * P_diff +
+#         coeff[12] * P_diff**3 + coeff[13] * h_diff * P_diff**3 +
+#         coeff[14] * h_diff**3 * P_diff**2 + coeff[15] * h_diff**2 * P_diff**3
+#     )
+
+#     return result
+
+# @partial(jit, static_argnums=(5, 6))
+def bicubic_interpolant(h, P, h_vals, P_vals, coeffs, Nh, Np, hmin, hmax, Lmin, Lmax):
+    """
+    Evaluate the bicubic interpolant at (h, P) using precomputed coefficients.
+    """
+    # Log-transform P
+    L = jnp.log(P)
+
+    # Normalize positions to [0, 1] cell coordinates
+    ii = ((h - hmin) / (hmax - hmin) * (Nh - 1))
+    i = ii.astype(int)
+    x = ii - i
+
+    jj = ((L - Lmin) / (Lmax - Lmin) * (Np - 1))
+    j = jj.astype(int)
+    y = jj - j
+
+
+    # Evaluate bicubic polynomial
+    result = jnp.zeros_like(h)  # use h shape
+    x_pow = jnp.ones_like(h)    # x^0
+
+    for m in range(4):  # m = x power
+        y_pow = jnp.ones_like(h)  # y^0 initially
+        for n in range(4):  # n = y power
+            c = coeffs[i, j, 4 * n + m]
+            result += c * x_pow * y_pow
+            y_pow = y_pow * y
+        x_pow = x_pow * x
+
+
+    return result
+
+
+
+# @jax.jit
+def inverse_interpolant_scalar_hD(h, D):
+    #Find the real(float) index
+    ii=((h-hmin)/(hmax-hmin)*(N-1))
+    #The integer part is the cell index
+    i=ii.astype(int)
+    #The remainder (for numerical stability better to use the difference)
+    #is instead the position within our interpolation cell.
+    x=ii-i
+    #find interval that contains the solution
+    xth=jnp.ones_like(h) #initialize x to the 0th power
+    #First we compute the nodal values, that is the values of D(h,P) where
+    #h is the actual enthalpy and P are grid values.
+    #TODO: instead of computing all the nodal values and then use sortedsearch
+    #to find the correct interval, we could do a binary search. This would
+    #constraint M to be a power of 2.
+    #Possible example (to be refined) to compute the node. Start with the node
+    #corresponding to j=M/2, then compute new index j=j+M/4*(2*(Dj>D)-1)
+    #then j=j+M/8*(2*(Dj>D)-1) and so on ...
+    #after log2(M) iteration we converged to the index j.
+    D_nodal=jnp.zeros(M)
+    for m in range(4):
+        D_nodal+=bicubic_coefficients[i,:,m]*xth
+        xth=xth*x
+    #We search more efficiently in which interval we have the solution
+    #if we assume a sorted vector.
+    #TODO: This assumes that P has a monotonic trend with respect to D
+    #at fixed h. This causes some problems and needs further investigation
+    if iD==cp.iSmass:
+        j=jax.numpy.searchsorted(-D_nodal,-D).astype(int)-1
+    else:
+        j=jax.numpy.searchsorted(D_nodal,D).astype(int)-1
+
+    #After we are in the unit square, that is for known i and j
+    #compute 1D cubic coefficients (as complex numbers to avoid promotion)
+    #Each coefficient is bj=sum(aij*x**i)
+    #Leading to the equation D=b0 + b1*y + b2*y**2 + b3*y**3
+    xth=jnp.ones_like(h)
+    b0 =jnp.zeros_like(h,dtype=complex128)
+    b1 =jnp.zeros_like(h,dtype=complex128)
+    b2 =jnp.zeros_like(h,dtype=complex128)
+    b3 =jnp.zeros_like(h,dtype=complex128)
+    for m in range(4):
+        b0 +=bicubic_coefficients[i,j,4*0+m]*xth
+        b1 +=bicubic_coefficients[i,j,4*1+m]*xth
+        b2 +=bicubic_coefficients[i,j,4*2+m]*xth
+        b3 +=bicubic_coefficients[i,j,4*3+m]*xth
+        xth=xth*x
+    #solve cubic equation - all three solutions
+    #TODO: if necessary, add solution for degenerate (quadratic and linear)
+    #For more information:https://en.wikipedia.org/wiki/Cubic_equation#General_cubic_formula
+    D0=b2*b2-3*b3*b1
+    D1=2*b2*b2*b2-9*b3*b2*b1+27*b3*b3*(b0-D)
+    C=((D1+(D1*D1-4*D0*D0*D0)**0.5)/2)**(1/3)
+    D0C=jax.lax.select(C==(0+0j),0+0j,D0/C)
+    z=jnp.array([1,-0.5+0.8660254037844386j,-0.5-0.8660254037844386j])
+    y=-1/(3*b3)*(b2+C*z+D0C/z)
+    #To find our solution we have two criteria:
+    #   -0 imaginary part
+    #   -real part between 0 and 1, that are the bounds of our cell
+    # We define a "badness" as the deviation from these critera, and pick the
+    # solution with the lowest badness
+    badness=jax.nn.relu(4*(jnp.real(y)-0.5)**2-1)+jnp.imag(y)**2
+    yreal=jnp.real(y[jnp.argmin(badness)])
+    jj=j+yreal
+    L=Lmin+jj*(Lmax-Lmin)/(M-1)
+    P=jnp.exp(L)
+    return P
+
+# @jax.jit
+def inverse_interpolant_scalar_DP(D, P):
+    # Convert pressure to log space
+    L = jnp.log(P)
+    
+    # Compute index along pressure grid
+    jj = ((L - Lmin) / (Lmax - Lmin) * (M - 1))
+    j = jj.astype(int)
+    y = jj - j  # fractional position in pressure direction
+
+    # Compute nodal D(h) values at fixed pressure (we'll search h index now)
+    yth = jnp.ones_like(D)
+    D_nodal = jnp.zeros(N)
+    for m in range(4):
+        D_nodal += bicubic_coefficients[:, j, m] * yth
+        yth = yth * y
+
+    # Search h-direction to find which cell to use
+    if iD == cp.iSmass:
+        i = jnp.searchsorted(-D_nodal, -D).astype(int) - 1
+    else:
+        i = jnp.searchsorted(D_nodal, D).astype(int) - 1
+
+    # Now build 1D cubic in x (h-direction) at fixed j
+    yth = jnp.ones_like(D)
+    b0 = jnp.zeros_like(D, dtype=complex128)
+    b1 = jnp.zeros_like(D, dtype=complex128)
+    b2 = jnp.zeros_like(D, dtype=complex128)
+    b3 = jnp.zeros_like(D, dtype=complex128)
+    for m in range(4):
+        b0 += bicubic_coefficients[i, j, m + 4*0] * yth
+        b1 += bicubic_coefficients[i, j, m + 4*1] * yth
+        b2 += bicubic_coefficients[i, j, m + 4*2] * yth
+        b3 += bicubic_coefficients[i, j, m + 4*3] * yth
+        yth = yth * y
+
+    # Solve cubic: D = b0 + b1*x + b2*x^2 + b3*x^3
+    D0 = b2*b2 - 3*b3*b1
+    D1 = 2*b2**3 - 9*b3*b2*b1 + 27*b3**2*(b0 - D)
+    C = ((D1 + jnp.sqrt(D1**2 - 4*D0**3)) / 2)**(1/3)
+    D0C = jax.lax.select(C == 0, 0 + 0j, D0 / C)
+    z = jnp.array([1, -0.5 + 0.8660254037844386j, -0.5 - 0.8660254037844386j])
+    x = -1/(3*b3)*(b2 + C*z + D0C/z)
+
+    # Pick root with lowest badness
+    badness = jax.nn.relu(4*(jnp.real(x)-0.5)**2 - 1) + jnp.imag(x)**2
+    xreal = jnp.real(x[jnp.argmin(badness)])
+
+    # Final result: compute h from i + x
+    ii = i + xreal
+    h = hmin + ii * (hmax - hmin) / (N - 1)
+    return h

jaxprop-0.3.1/jaxprop/components/__init__.py

@@ -0,0 +1,3 @@
+# from ..jax_helpers import *
+from .nozzle_model_core import *
+from .nozzle_model_solver import *