TfbiChunkSolver

class PlasmaCalcs.addons.tfbi.tfbi_solver.TfbiChunkSolver(cc, dst, ions=None, *, kres='low', mod=UNSET, lmod=UNSET, ang=UNSET, tfbi_all=True, drel_cls=None, pre_hook=None, post_hook=None)

Bases: TfbiSolver

high-level interface for solving TFBI across many physical parameters, using chunks.

Chunking details are determined by cc.chunks; see help(cc.chunker) for more info.

Aside from the chunking, this works very similarly to TfbiSolver.

Call TfbiChunkSolver to solve TFBI.

Example:

import PlasmaCalcs as pc
cc = pc.PlasmaCalculator(…)   # <– your PlasmaCalculator of choice
cc.chunks = dict(…)  # <– your chunking parameters of choice
solver = pc.TfbiChunkSolver(cc)
solution = solver()             # alias: solver.solve()

solution is an xarray.Dataset with all relevant quantities (see TfbiLoader.get_tfbi_all),

and ‘omega’ telling roots with largest imaginary part.

One good choice of chunks is to subsample spatial dimensions, like:

XSUBSAMPLE, XBIGSTEP = 64, 256
XSLICES = [slice(ix0, None, XBIGSTEP) for ix0 in range(0, XBIGSTEP, XSUBSAMPLE)]
cc.chunks = dict(x_slices=XSLICES)
# this will do 4 chunks total, each chunk handling every 256th point along x.
# so, every chunk gives the “full” picture, though very low resolution.
# but, the chunks are offset from each other by 64, so altogether
# they actual combine to give every 64th point along x.
# then, one can rerun later with smaller XSUBSAMPLE (e.g. 32)
# to improve resolution while re-using existing results!!!

For more precise control of solving process, consider:

(copied from TfbiSolver docs): the following pattern & notes:

import tfbi_theory as tt
import PlasmaCalcs as pc
cc = … # any PlasmaCalculator object from PlasmaCalcs.
ds0 = cc.tfbi_ds()
kp = tt.kPickerLowres(ds0)
dsk = kp.get_ds()   # copy of ds0, but with ds[‘k’] = k from kPicker.
drel = tt.TfbiDisprelC.from_ds(dsk)
dsR = drel.solve()  # copy of dsk, but with ds[‘omega’] = solution to TFBI theory!

Notes (copied from TfbiSolver docs):

if cc has more than ~4 ions, you will want to drop some or group them somehow.

e.g. for MhdMultifluidCalculator, before calling tfbi_ds():

cc.use_mix_heavy_ions(0.3, m_mean_mode=’density’)

You can also pick ions directly during cc.tfbi_ds().

e.g. for Bifrost chromosphere analysis, where n[He_II] < 1e-6 * ne, I use:

cc.tfbi_ds(ions=[i for i in cc.fluids if i.q==1 and i!=’He_II’])

After finishing solving, you might want to save the result,

e.g. dsR.pc.save(‘filename’) saves result to ‘filename.pcxarr’;

can load it later via pc.xarray_load(‘filename.pcxarr’).

If you are solving across “many” points (e.g., more than ~5000),

you might want to use chunking instead. see TfbiChunkSolver;
also note cc.tfbi_solver(…) automatically uses TfbiChunkSolver
if appropriate, i.e. if cc.chunks is set (where cc is a TfbiLoader).

BUT, the pattern above doesn’t include chunking.

The FULL pattern, with chunking included, looks roughly like:

slicers = cc.chunker().slicers()

locs = [os.path.join(dst, cc.title_with_slices(s)+’.pcxarr’) for s in slicers]

todo = [(s, loc) for (s, loc) in zip(slicers, locs) if not os.path.exists(loc)]

for s, p in todo:

cc.slices = s
cc.chunks = None
# (*1) (point of interest, referenced below)
dsR = … # <– copy-paste most of the pattern from above!
dsR.pc.save(loc)
# (*2) (point of interest, referenced below)

[TODO][EFF] would be slightly more efficient to avoid recomputing casc for each chunk.

However, for “reasonable chunk sizes”, this is only a ~10% speedup.

INPUTS – used here, not just passed directly to super().

cc: PlasmaCalculator

PlasmaCalculator object used to load the data.

Should be a TfbiLoader subclass. (PlasmaCalculator satisfies this by default,

assuming successful import SymSolver and import tfbi_theory.)

cc.chunks MUST be nonempty; otherwise, use TfbiSolver instead.

dst: str

directory where chunk results should be saved. Internally stored as abspath.

pre_hook: None or callable of cc

if provided, called as cc = pre_hook(cc) at point (*1) in the pattern above,

i.e. before solving each chunk.

E.g. def pre_hook(cc): cc.tfbi_mask(); return cc

would apply tfbi_mask() to each chunk before solving.
(Might want this because slicing after masking is hard,
but masking after slicing is easy, in present implementation.)

E.g., to artificially set E_un0_perpmod_B to test constant E/B speed across box:

def pre_hook(cc):

result = cc.copy()
EBspeed_grid = pc.xr1d([1000, 3000, 5000, 7000], ‘EBspeed’)
result.set(‘E_un0_perpmod_B’, cc(‘mod_B’) * EBspeed_grid)
return result

post_hook: None or callable of (cc, dsR, loc)

if provided, called as post_hook(cc, dsR, loc) at point (*2) in the pattern above,

i.e. after solving each chunk and saving the result.

E.g., to save a growth rate plot in each chunk’s loc directory:

def post_hook(cc, dsR, loc):

dsR.pc.unmask_var(‘omega’).it.growth_kmax().it.growthplot()
plt.savefig(os.path.join(loc, ‘growthplot’))
plt.close()

INPUTS – passed directly to super(), i.e. works the same way as TfbiSolver.

ions: None or specifier of multiple fluids (e.g. slice, or list of strs)

None –> use cc.fluids.ions()

ions are determined when called, not during __init__.

print warning if this specifies more than DEFAULTS.ADDONS.TFBI_MAX_NUM_IONS ions

(default: 5), because then solving will be slow and may be inaccurate.

kres: ‘low’, ‘mid’, or ‘high’

resolution in k-space. Tells which self.kPicker_cls to use.
‘low’ –> tfbi_theory.kPickerLowres. Recommended if solving across many (e.g. >1000) points.
‘mid’ –> tfbi_theory.kPickerMidres. Recommended if solving across a few (e.g. 10 to 100) points.
‘high’ –> tfbi_theory.kPickerHighres. Recommended if solving at only 1 point.

mod, lmod, ang: UNSET or dict

passed directly to kPicker if provided. Can specify k values other than the defaults.

see help(self.kPicker_cls) for more details.

tfbi_all: bool

whether to compute all relevant tfbi vars, ds0 = cc(‘tfbi_all’).

False –> compute only the necessary vars, ds0 = cc(‘tfbi_inputs’).

drel_cls: None, str, or class

tfbi_theory class to use for solving TFBI theory.
None –> use self.drel_cls default: tt.TfbiDisprelC
str –> use getattr(tt, drel_cls) to get the class.

Methods

`chunk_filepath`(slicer)	return filepath to use for results from chunk defined by slicer.
`done`()	list of (slicer, filepath) tuples for all chunks which have already been solved.
`n_chunks`()	total number of chunks implied by self.cc.chunks.
`slicers_and_filepaths`()	return list of (slicer, filepath) tuples for all chunks implied by self.cc.chunks.
`solve`(*[, verbose, mergeload])	solve TFBI theory using chunking.
`todo`()	list of (slicer, filepath) tuples for all chunks which have NOT yet been solved.
`_default_drel_cls`()	default drel_cls.
`_repr_contents`()	return list of strings for contents inside of __repr__
`_warn_if_too_many_ions`()	print warning if self.ions specifies too many ions.

Attributes

`drel_cls`	class to use for solving TFBI theory.
`drel_cls_name`
`elastic`	whether this solver assumes elastic collisions for all species
`ions_explicit`	list of ions from self.cc which would be used during self.cc.tfbi_ds()
`kPickerHighres_cls_name`
`kPickerLowres_cls_name`
`kPickerMidres_cls_name`
`kPicker_cls`	kPicker class to use for choosing wavevectors to consider.
`solution`	result of self.solve(); alias to self.dsR.
`solved`	tells number of chunks solved and total, as "Nsolved_of_Ntotal" string.

_default_drel_cls(): default drel_cls. getattr(tt, self.drel_cls_name)

_repr_contents(): return list of strings for contents inside of __repr__

_warn_if_too_many_ions(): print warning if self.ions specifies too many ions.

chunk_filepath(slicer): return filepath to use for results from chunk defined by slicer.

Slicer should be a dict of slices, appropriate as cc.slices = slicer.

done(): list of (slicer, filepath) tuples for all chunks which have already been solved.

(To check if solved, just checks whether filepath exists!)

property drel_cls: class to use for solving TFBI theory. Default: tt.TfbiDisprelC.

property elastic: whether this solver assumes elastic collisions for all species

property ions_explicit: list of ions from self.cc which would be used during self.cc.tfbi_ds()

property kPicker_cls: kPicker class to use for choosing wavevectors to consider. Depends on self.kres:

‘low’ –> tfbi_theory.kPickerLowres

‘mid’ –> tfbi_theory.kPickerMidres

‘high’ –> tfbi_theory.kPickerHighres

n_chunks(): total number of chunks implied by self.cc.chunks.

slicers_and_filepaths(): return list of (slicer, filepath) tuples for all chunks implied by self.cc.chunks.

property solution: result of self.solve(); alias to self.dsR.

solve(*, verbose=2, mergeload=True, **kw_growth_root): solve TFBI theory using chunking.

See help(type(self)) for more details.

Writes results to files, as determined by self.dst.

Never overwrite any files; assumes pre-existing files are correct solutions.

verbose: bool or int

whether to print progress updates (highly recommended).

0: print nothing

>=1: print updates about which chunk is being done.

>=2: print tfbi_solver updates within each chunk.

(chunks might each take a few minutes)

additional kwargs get passed directly to drel.solve().

(options include: ncpu, ncoarse, careful)

mergeload: bool

whether to return full solution, mergeloaded across all chunks.

True –> return full solution.

False –> return filepath to where chunks are stored (i.e., self.dst).

Either way, the full solution can later be loaded via pc.xarray_mergeload(dst).

property solved: tells number of chunks solved and total, as “Nsolved_of_Ntotal” string.

E.g., “3_of_10” if 3 out of 10 chunks have existing solutions.

todo(): list of (slicer, filepath) tuples for all chunks which have NOT yet been solved.

(To check if solved, just checks whether filepath exists!)