re-strucuring basis

Author: jngaravitoc

EXPtools/basis/__init__.py

@@ -0,0 +1,3 @@
+from .profiles import Profiles
+from .basis_utils import *
+from .makemodel import *

EXPtools/basis/basis_utils.py

@@ -0,0 +1,218 @@
+import os
+import re
+import yaml
+import numpy as np
+import pyEXP
+from EXPtools.basis.makemodel import make_model
+
+def load_basis(config_name, cache_dir=None):
+    """
+    Load a basis configuration from a YAML file and initialize a Basis object.
+
+    Parameters
+    ----------
+    config_name : str
+        Path to the YAML configuration file. If the provided filename does not 
+        end with `.yaml`, the extension is automatically appended.
+    cache_dir : str (optional)
+        Path to modelname and cachename (assumes both are in the same directory)
+        if None assumes it is the current working directory.
+
+    Returns
+    -------
+    basis : pyEXP.basis.Basis
+        An initialized Basis object created from the configuration.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the specified YAML file does not exist.
+    """
+
+    # Check file existence
+    if not os.path.exists(config_name):
+        raise FileNotFoundError(f"Configuration file not found: {config_name}")
+
+    # Load YAML safely
+    with open(config_name, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    if cache_dir:
+        config = re.sub(r"(modelname:\s*)(\S+)", rf"\1{cache_dir}\2", config)
+        config = re.sub(r"(cachename:\s*)(\S+)", rf"\1{cache_dir}\2", config)
+    # Build basis from configuration
+    basis = pyEXP.basis.Basis.factory(config)
+    return basis
+    
+def check_basis_params(basis_id, **kwargs):
+    """
+    Check that the required keyword arguments for a given basis are provided.
+
+    Parameters
+    ----------
+    basis_id : str
+        The identifier of the basis set. 
+        Accepted values are:
+        - ``'sphereSL'`` : Spherical basis with Stäckel-like mapping.
+        - ``'cylinder'`` : Cylindrical basis.
+    **kwargs : dict
+        Arbitrary keyword arguments corresponding to the basis parameters.
+        The required keys depend on ``basis_id``:
+
+        - For ``'sphereSL'``:
+          ['Lmax', 'mmax', 'modelname', 'rmapping', 'cachename']
+
+        - For ``'cylinder'``:
+          ['acyl', 'hcyl', 'nmaxfid', 'lmaxfid', 'mmax', 'nmax', 'ncylodd',
+           'ncylnx', 'ncylny', 'rnum', 'pmun', 'tnum', 'vflag', 'logr', 'cachename']
+
+    Returns
+    -------
+    bool
+        Returns ``True`` if all mandatory parameters are present.
+
+    Raises
+    ------
+    KeyError
+        If one or more mandatory keyword arguments are missing for the selected basis.
+    AttributeError
+        If ``basis_id`` is not recognized (must be either 'sphereSL' or 'cylinder').
+
+    Examples
+    --------
+    check_basis_params('sphereSL', Lmax=4, nmax=4, modelname='hernquist.txt',
+    ...                    rmapping=1, cachename='cache_hernquist.txt')
+    True
+
+    check_basis_params('cylinder', acyl=1.0, hcyl=2.0)  
+    Traceback (most recent call last):
+    ...
+    KeyError: "Missing mandatory keyword arguments missing: [...]"
+    """
+    
+    if basis_id == 'sphereSL':
+        mandatory_keys = ['Lmax', 'nmax', 'modelname', 'rmapping', 'cachename']
+        missing = [key for key in mandatory_keys if key not in kwargs]
+        if missing:
+            raise KeyError(f"Missing mandatory keyword arguments missing: {missing}")
+        return True
+    elif basis_id == 'cylinder':
+        mandatory_keys = ['acyl', 'hcyl', 'nmaxfid', 'lmaxfid', 
+                           'mmax', 'nmax', 'ncylodd', 'ncylnx', 
+                           'ncylny', 'rnum', 'pnum', 'tnum', 'vflag', 'logr', 'cachename']  
+        missing = [key for key in mandatory_keys if key not in kwargs]
+        if missing:
+            raise KeyError(f"Missing mandatory keyword arguments missing: {missing}")
+        return True
+    else: 
+        raise AttributeError(f"basis id {basis_id} not found. Please chose between sphereSL or cylinder")
+
+
+def write_config(basis_id, float_fmt_rmin="{:.7f}", float_fmt_rmax="{:.3f}",
+                float_fmt_rmapping="{:.3f}", **params):
+    """
+    Create a YAML configuration file string for building a basis model.
+
+    Parameters
+    ----------
+    basis_id : str
+        Identifier of the basis model. Must be either 'sphereSL' or 'cylinder'.
+    float_fmt_rmin : str, optional
+        Format string for rmin (default ``"{:.7f}"``).
+    float_fmt_rmax : str, optional
+        Format string for rmax (default ``"{:.3f}"``).
+    float_fmt_rmapping : str, optional
+        Format string for rmapping (default ``"{:.3f}"``).
+    **params : dict
+        Additional keyword arguments required depending on the basis type:
+
+        - For ``sphereSL``:
+          ['lmax', 'nmax', 'rmapping', 'modelname', 'cachename']
+
+        - For ``cylinder``:
+          ['acyl', 'hcyl', 'nmaxfid', 'lmaxfid', 'mmax', 'nmax',
+           'ncylodd', 'ncylnx', 'ncylny', 'rnum', 'pnum', 'tnum',
+           'vflag', 'logr', 'cachename']
+
+    Returns
+    -------
+    str
+        YAML configuration file contents.
+
+    Raises
+    ------
+    KeyError
+        If mandatory parameters for the given basis are missing.
+    FileNotFoundError
+        If ``modelname`` is required but cannot be opened.
+    ValueError
+        If the model file does not contain valid radius data.
+    """
+
+    check_basis_params(basis_id, **params)
+
+   
+    if basis_id == "sphereSL":
+        modelname = params["modelname"]
+        try:
+            R = np.loadtxt(modelname, skiprows=3, usecols=0)
+        except OSError as e:
+            raise FileNotFoundError(f"Could not open model file '{modelname}'") from e
+        if R.size == 0:
+            raise ValueError(f"Model file '{modelname}' contains no radius data")
+
+        rmin, rmax, numr = R[0], R[-1], len(R)
+        params["rmin"] = rmin
+        params["rmax"] = rmax
+        params["numr"] = numr
+
+
+    config_dict = {
+        "id": basis_id,
+        "parameters": params
+        }
+    
+    
+    return yaml.dump(config_dict, sort_keys=False)
+
+def make_basis(R, D, Mtotal, basis_params, modelname="test_model.txt", cachename='test_cache.txt'):
+    """
+    Construct a basis from a given radial density profile.
+
+    Parameters
+    ----------
+    R : array_like
+        Radial grid points (e.g., radii at which density `D` is defined).
+    D : array_like
+        Density values corresponding to each radius in `R`.
+    Mtotal : float, optional
+        Total mass normalization (default is 1.0).
+    basis_params : dict 
+        basis parameters e.g., basis_id, nmax, lmax
+
+    Returns
+    -------
+    basis : pyEXP.basis.Basis
+        A basis object initialized with the given density model.
+
+    Notes
+    -----
+    - This function wraps `makemodel.makemodel` to generate a model from 
+      the supplied density profile and total mass.
+    - It then builds a basis either spherical (`sphereSL`) or cylindrical using `EXPtools.make_config`
+      and returns the corresponding `pyEXP` basis object.
+    """
+
+    R, D, _, _ = make_model(
+        R, D, Mtotal=Mtotal, 
+        output_filename=modelname
+    )
+
+    config = write_config(
+        basis_id=basis_params['basis_id'],
+        modelname=modelname,
+        cachename=cachename
+    )
+
+    basis = pyEXP.basis.Basis.factory(config)
+    return basis

EXPtools/basis/makemodel.py

@@ -0,0 +1,135 @@
+import numpy as np
+
+def write_table(tablename, radius, density, mass, potential, fmt="%.6e"):
+    """
+    Write a table of radius, density, mass, and potential values to a text file.
+
+    Parameters
+    ----------
+    tablename : str
+        Output filename.
+    radius, density, mass, potential : array-like
+        Arrays of physical quantities, all with the same length.
+    fmt : str, optional
+        Format string for numerical values. Defaults to scientific notation with 6 decimals.
+
+    Notes
+    -----
+    Writes the table in the following format:
+        ! <tablename>
+        ! R    D    M    P
+        <Nrows>
+        <radius> <density> <mass> <potential>
+    """
+    # Convert inputs to NumPy arrays (for safety and performance)
+    radius = np.asarray(radius)
+    density = np.asarray(density)
+    mass = np.asarray(mass)
+    potential = np.asarray(potential)
+
+    # Stack data into a single 2D array for fast writing
+    data = np.column_stack((radius, density, mass, potential))
+
+    header = f"! {tablename}\n! R    D    M    P\n{len(radius)}"
+    np.savetxt(tablename, data, fmt=fmt, header=header, comments="")
+
+def make_model(radius, density, Mtotal, output_filename='', physical_units=False, verbose=True):
+    """
+    Generate an EXP-compatible spherical basis function table.
+
+    Parameters
+    ----------
+    radius : array-like
+        Radii at which the density values are evaluated.
+    density : array-like
+        Density values corresponding to radius.
+    Mtotal : float
+        Total mass of the model, used for normalization.
+    output_filename : str, optional
+        Name of the output file to save the table. If empty, no file is written.
+    physical_units : bool, optional
+        If True, disables scaling and returns physical values (default: False).
+    verbose : bool, optional
+        If True, prints scaling information.
+
+    Returns
+    -------
+    result : dict
+        Dictionary with the following keys:
+        - 'radius' : ndarray
+            Scaled radius values.
+        - 'density' : ndarray
+            Scaled density values.
+        - 'mass' : ndarray
+            Scaled enclosed mass values.
+        - 'potential' : ndarray
+            Scaled potential values.
+    
+    Modified from Mike Petersen's original code:
+    https://gist.github.com/michael-petersen/ec4f20641eedac8f63ec409c9cc65ed7
+    """
+    
+    EPS_MASS = 1e-15
+    EPS_R = 1e-10
+    
+    Rmax = np.nanmax(radius)
+    
+    mass = np.zeros_like(density)
+    pwvals = np.zeros_like(density)
+
+    mass[0] = 1.e-15
+    pwvals[0] = 0.
+
+    dr = np.diff(radius)  
+
+    # Midpoint integration for enclosed mass and potential
+    mass_contrib = 2.0 * np.pi * (
+        radius[:-1]**2 * density[:-1] + radius[1:]**2 * density[1:]
+    ) * dr
+
+    pwvals_contrib = 2.0 * np.pi * (
+        radius[:-1] * density[:-1] + radius[1:] * density[1:]
+    ) * dr
+
+    # Now cumulative sum to get the arrays
+    mass = np.concatenate(([EPS_MASS], EPS_MASS + np.cumsum(mass_contrib)))
+    pwvals = np.concatenate(([0.0], np.cumsum(pwvals_contrib)))
+    
+    potential = -mass / (radius + EPS_R) - (pwvals[-1] - pwvals)
+
+    M0 = mass[-1]
+    R0 = radius[-1]
+
+    Beta = (Mtotal / M0) * (R0 / Rmax)
+    Gamma = np.sqrt((M0 * R0) / (Mtotal * Rmax)) * (R0 / Rmax)
+
+    if verbose:
+        print(f"! Scaling: R = {Rmax}  M = {Mtotal}")
+
+    rfac = Beta**-0.25 * Gamma**-0.5
+    dfac = Beta**1.5 * Gamma
+    mfac = Beta**0.75 * Gamma**-0.5
+    pfac = Beta
+
+    if physical_units:
+        rfac = dfac = mfac = pfac = 1.0
+
+    if verbose:
+        print(f"Scaling factors: rfac = {rfac}, dfac = {dfac}, mfac = {mfac}, pfac = {pfac}")
+
+    if output_filename:
+        write_table(
+            output_filename,
+            radius * rfac,
+            density * dfac,
+            mass * mfac,
+            potential * pfac
+        )
+
+    return {
+        "radius": radius * rfac,
+        "density": density * dfac,
+        "mass": mass * mfac,
+        "potential": potential * pfac,
+    }
+