lightweaver 0.15.0__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

lightweaver/barklem.py

@@ -0,0 +1,151 @@
+from typing import TYPE_CHECKING, Sequence, Tuple
+
+import numpy as np
+from scipy.interpolate import RectBivariateSpline
+from scipy.special import gamma
+
+import lightweaver.constants as Const
+from .atomic_table import PeriodicTable
+from .utils import get_data_path
+
+if TYPE_CHECKING:
+    from .atomic_model import AtomicLine, AtomicModel
+
+DeltaNeff = 0.1
+
+class BarklemTable:
+    '''
+    Storage for each table of Barklem data for Van der Waals approximation.
+    '''
+    def __init__(self, path: str, neff0: Tuple[float, float]):
+        data = np.genfromtxt(path, comments='c')
+        shape = data.shape
+        self.cross = data[:shape[0]//2]
+        self.alpha = data[shape[0]//2:]
+
+        self.neff1 = neff0[0] + np.arange(shape[0]//2) * DeltaNeff
+        self.neff2 = neff0[1] + np.arange(shape[1]) * DeltaNeff
+
+class BarklemCrossSectionError(Exception):
+    '''
+    Raised if the Barklem cross-section cannot be applied to the atom in
+    question.
+    '''
+    pass
+
+class Barklem:
+    '''
+    Storage for all three Barklem cross-section cases and application via the
+    `get_active_cross_section` function.
+    '''
+    barklem_sp = BarklemTable(get_data_path() + 'Barklem_spdata.dat', (1.0, 1.3))
+    barklem_pd = BarklemTable(get_data_path() + 'Barklem_pddata.dat', (1.3, 2.3))
+    barklem_df = BarklemTable(get_data_path() + 'Barklem_dfdata.dat', (2.3, 3.3))
+
+    @classmethod
+    def get_active_cross_section(cls, atom: 'AtomicModel',
+                                 line: 'AtomicLine',
+                                 vals: Sequence[float]) -> Sequence[float]:
+        '''
+        Returns the cross section data for use in the Van der Waals collisional
+        broadening routines.
+        See:
+
+          - Anstee & O'Mara 1995, MNRAS 276, 859-866
+
+          - Barklem & O'Mara 1998, MNRAS 300, 863-871
+
+          - Unsold:
+
+            - Traving 1960, "Uber die Theorie der Druckverbreiterung
+              von Spektrallinien", p 91-97
+
+            - Mihalas 1978, p. 282ff, and Table 9-1/
+
+        Returns
+        -------
+        result : list of 3 float
+            Barklem cross-section, Barklem alpha, Helium contribution
+            following Unsold (always 1.0)
+        '''
+        i = line.i
+        j = line.j
+
+        SOrbit = 0
+        POrbit = 1
+        DOrbit = 2
+        FOrbit = 3
+
+        result = [vals[0], vals[1], 0.0]
+
+        # Follows original RH version. Interpolate tables if sigma < 20.0, otherwise
+        # assume the provided values are the coefficients
+        if vals[0] < 20.0:
+            if atom.levels[i].stage > 0:
+                raise BarklemCrossSectionError('Atom is not neutral.')
+
+            # Find principal quantum numbers
+            # try:
+            #     lowerNum = determinate(atom.levels[i])
+            #     upperNum = determinate(atom.levels[j])
+            # except CompositeLevelError:
+            #     raise BarklemCrossSectionError()
+            lowerNum = atom.levels[i].L
+            upperNum = atom.levels[j].L
+            if lowerNum is None or upperNum is None:
+                raise BarklemCrossSectionError('L not provided for levels.')
+
+            nums = (lowerNum, upperNum)
+
+            # Check is a Barklem case applies
+            if nums == (SOrbit, POrbit) or nums == (POrbit, SOrbit):
+                table = cls.barklem_sp
+            elif nums == (POrbit, DOrbit) or nums == (DOrbit, POrbit):
+                table = cls.barklem_pd
+            elif nums == (DOrbit, FOrbit) or nums == (FOrbit, DOrbit):
+                table = cls.barklem_df
+            else:
+                raise BarklemCrossSectionError('Not a valid shell combination.')
+
+            Z = atom.levels[j].stage + 1
+            # Find index of continuum level
+            ic = j + 1
+            while atom.levels[ic].stage < atom.levels[j].stage + 1:
+                ic += 1
+
+            deltaEi = (atom.levels[ic].E - atom.levels[i].E) * Const.HC_CM
+            deltaEj = (atom.levels[ic].E - atom.levels[j].E) * Const.HC_CM
+            E_Rydberg = Const.ERydberg / (1.0 + Const.MElectron
+                                           / (atom.element.mass * Const.Amu))
+
+            neff1 = Z * np.sqrt(E_Rydberg / deltaEi)
+            neff2 = Z * np.sqrt(E_Rydberg / deltaEj)
+
+            if nums[0] > nums[1]:
+                neff1, neff2 = neff2, neff1
+
+            if not (table.neff1[0] <= neff1 <= table.neff1[-1]):
+                raise BarklemCrossSectionError('neff1 outside table.')
+            if not (table.neff2[0] <= neff2 <= table.neff2[-1]):
+                raise BarklemCrossSectionError('neff2 outside table.')
+
+
+            result[0] = float(RectBivariateSpline(table.neff1, table.neff2,
+                                                  table.cross)(neff1, neff2))
+            result[1] = float(RectBivariateSpline(table.neff1, table.neff2,
+                                                  table.alpha)(neff1, neff2))
+
+        reducedMass = Const.Amu / (1.0 / PeriodicTable[1].mass + 1.0 / atom.element.mass)
+        meanVel = np.sqrt(8.0 * Const.KBoltzmann / (np.pi * reducedMass))
+        sigma = result[0]
+        alpha = result[1]
+        crossSection = sigma * Const.RBohr**2 * (meanVel / 1.0e4)**(-alpha)
+
+        # NOTE(cmo): This is w/N/T^(0.5*(1.0-alpha)) (eq 3 without temperature contrib, multiplied by 2 for half-width)
+        result[0] = 2.0 * ((4.0 / np.pi)**(alpha / 2.0)
+                     * gamma(2.0 - alpha / 2.0) * meanVel * crossSection)
+
+        # Use Unsold for Helium contribution
+        result[2] = 1.0
+
+        return result

lightweaver/benchmark.py

@@ -0,0 +1,113 @@
+import time
+
+import numpy as np
+from tqdm import tqdm
+from weno4 import weno4
+
+from lightweaver.config import update_config_file
+from .atmosphere import Atmosphere, ScaleType
+from .atomic_set import RadiativeSet
+from .config import params as rcParams
+from .config import get_home_config_path, update_config_file
+from .fal import Falc82
+from .rh_atoms import CaII_atom, H_6_atom
+from .simd_management import get_available_simd_suffixes
+from .LwCompiled import LwContext
+
+__all__ = ['benchmark']
+
+def configure_context(Nspace=500, fsIterScheme=None):
+    '''
+    Configure a FALC context (with more or fewer depth points), 1 thread and a
+    particular iteration scheme. For use in benchmarking.
+
+    Parameters
+    ----------
+    Nspace : int, optional
+        Number of spatial points to interpolate the atmosphere to. (Default: 500)
+    fsIterScheme : str, optional
+        The fsIterScheme to use in the Context. (Default: None, i.e. read from
+        the user's config)
+    '''
+    fal = Falc82()
+    interp = lambda x: weno4(np.linspace(0,1,Nspace), np.linspace(0,1,fal.Nspace), x)
+
+    atmos = Atmosphere.make_1d(ScaleType.Geometric, interp(fal.height),
+                               temperature=interp(fal.temperature), vlos=interp(fal.vlos),\
+                               vturb=interp(fal.vturb), ne=interp(fal.ne),
+                               nHTot=interp(fal.nHTot))
+    atmos.quadrature(5)
+    aSet = RadiativeSet([H_6_atom(), CaII_atom()])
+    aSet.set_active('H', 'Ca')
+    eqPops = aSet.compute_eq_pops(atmos)
+    spect = aSet.compute_wavelength_grid()
+    ctx = LwContext(atmos, spect, eqPops, fsIterScheme=fsIterScheme)
+    return ctx
+
+def benchmark(Niter=50, Nrep=3, verbose=True, writeConfig=True, warmUp=True):
+    '''
+    Benchmark the various SIMD implementations for Lightweaver's formal solver
+    and iteration functions.
+
+    Parameters
+    ----------
+    Niter : int, optional
+        The number of iterations to use for each scheme. (Default: 50)
+    Nrep : int, optional
+        The number of repetitions to average for each scheme. (Default: 3)
+    verbose : bool, optional
+        Whether to print information as the function runs. (Default: True)
+    writeConfig : bool, optional
+        Whether to writ the optimal method to the user's config file. (Default:
+        True)
+    warmUp : bool, optional
+        Whether to run a Context first (discarded) to ensure that all numba jit
+        code is jitted and warm. (Default: True)
+    '''
+    timer = time.perf_counter
+
+    if verbose:
+        print('This will take a couple of minutes...')
+
+    suffixes = get_available_simd_suffixes()
+    suffixes = ['scalar'] + suffixes
+    methods = [f'mali_full_precond_{suffix}' for suffix in suffixes]
+
+    if warmUp:
+        ctx = configure_context(fsIterScheme=methods[0])
+        for _ in range(max(Niter // 5, 10)):
+            ctx.formal_sol_gamma_matrices(printUpdate=False)
+
+    timings = [0.0] * len(suffixes)
+    it = tqdm(methods * Nrep) if verbose else methods * Nrep
+    for idx, method in enumerate(it):
+        ctx = configure_context(fsIterScheme=method)
+        start = timer()
+        for _ in range(Niter):
+            ctx.formal_sol_gamma_matrices(printUpdate=False)
+        end = timer()
+        duration = (end - start)
+        timings[idx % len(methods)] += duration
+
+    timings = [t / Nrep for t in timings]
+    if verbose:
+        for idx, method in enumerate(methods):
+            print(f'Timing for method "{method}": {timings[idx]:.3f} s '
+                  f'({Niter} iterations, {Nrep} repetitions)')
+
+    if writeConfig:
+        minTiming = min(timings)
+        minIdx = timings.index(minTiming)
+        if verbose:
+            print(f'Selecting method: {methods[minIdx]}')
+
+        impl = suffixes[minIdx]
+        rcParams['SimdImpl'] = impl
+
+        path = get_home_config_path()
+        if verbose:
+            print(f'Writing config to \'{path}\'...')
+        update_config_file(path)
+
+    if verbose:
+        print('Benchmark complete.')