Source code for lightweaver.nr_update

import numpy as np

from .atomic_set import lte_pops
from .atomic_table import PeriodicTable


[docs] def nr_post_update(self, fdCollisionRates=True, hOnly=False, timeDependentData=None, chunkSize=5, ngUpdate=None, printUpdate=None, extraParams=None): ''' Compute the Newton-Raphson terms for updating the electron density through charge conservation. Is attached to the Context object. Parameters ---------- fdCollisionRates : bool, optional Whether to use a finite difference approximation to the collisional rates to find the population gradient WRT ne (default: True i.e. use finite-difference, if False, collisional rates are ignored for this process.) hOnly : bool, optional Ignore atoms other than Hydrogen (the primary electron contributor) (default: False) timeDependentData : dict, optional The presence of this argument indicates that the time-dependent formalism should be used. Should contain the keys 'dt' with a floating point timestep, and 'nPrev' with a list of population vectors from the start of the time integration step, in the order of the active atoms. This latter term can be obtained from the previous state provied by `Context.time_dep_update`. chunkSize : int, optional Not currently used. ngUpdate : bool, optional Whether to apply Ng Acceleration (default: None, to apply automatic behaviour), will only accelerate if the counter on the Ng accelerator has seen enough steps since the previous acceleration (set in Context initialisation). printUpdate : bool, optional Whether to print information on the size of the update (default: None, to apply automatic behaviour). extraParams : dict, optional Dict of extra parameters to be converted through the `dict2ExtraParams` function and passed onto the C++ core. Returns ------- dPops : float The maximum relative change of any of the NLTE populations in the atmosphere. ''' if self.activeAtoms[0].element != PeriodicTable[1]: raise ValueError('Calling nr_post_update without Hydrogen active.') if ngUpdate is None: ngUpdate = self.conserveCharge if printUpdate is None: printUpdate = ngUpdate atoms = self.activeAtoms[:1] if hOnly else self.activeAtoms crswVal = self.crswCallback.val if hOnly: backgroundAtoms = [model for model in self.kwargs['spect'].radSet if model.element != PeriodicTable[1]] else: backgroundAtoms = self.kwargs['spect'].radSet.detailedAtoms + self.kwargs['spect'].radSet.passiveAtoms backgroundNe = np.zeros_like(self.atmos.ne) for atomModel in backgroundAtoms: lteStages = np.array([l.stage for l in atomModel.levels]) atom = self.kwargs['eqPops'].atomicPops[atomModel.element] backgroundNe += (lteStages[:, None] * atom.n[:, :]).sum(axis=0) neStart = np.copy(self.atmos.ne) dC = [] if fdCollisionRates: for atom in atoms: atom.compute_collisions(fillDiagonal=True) Cprev = np.copy(atom.C) pertSize = 1e-4 pert = neStart * pertSize self.atmos.ne[:] += pert nStarPrev = np.copy(atom.nStar) atom.nStar[:] = lte_pops(atom.atomicModel, self.atmos.temperature, self.atmos.ne, atom.nTotal) atom.compute_collisions(fillDiagonal=True) self.atmos.ne[:] = neStart atom.nStar[:] = nStarPrev dC.append(crswVal * (atom.C - Cprev) / pert) atom.C[:] = Cprev self._nr_post_update_impl(atoms, dC, backgroundNe, timeDependentData=timeDependentData, chunkSize=chunkSize, extraParams=extraParams) self.eqPops.update_lte_atoms_Hmin_pops(self.atmos.pyAtmos, conserveCharge=False, quiet=True) if ngUpdate: update = self.rel_diff_ng_accelerate(printUpdate=printUpdate) else: update = self.rel_diff_pops(printUpdate=printUpdate) neDiff = ((np.asarray(self.atmos.ne) - neStart) / np.asarray(self.atmos.ne)) neDiffMaxIdx = neDiff.argmax() neDiffMax = neDiff[neDiffMaxIdx] update.updatedNe = True update.dNeMax = neDiffMax update.dNeMaxIdx = neDiffMaxIdx return update