135 lines
3.5 KiB
Python
135 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from .distributions import DeltaDistribution, LogGaussianDistribution
|
|
from .motions import RandomJump, TetrahedralJump
|
|
from .parameter import *
|
|
|
|
|
|
def parse(config_file: str) -> Parameter:
|
|
with open(config_file, 'r') as f:
|
|
parameter: dict = json.load(f)
|
|
|
|
ste = _parse_ste(parameter.get('stimulated_echo'))
|
|
spec = _parse_spectrum(parameter.get('spectrum'))
|
|
|
|
if ste is None and spec is None:
|
|
raise ValueError("No parameter for STE or spectra given")
|
|
|
|
t_max = 0
|
|
if spec is not None:
|
|
t_max = max(spec.t_max, t_max)
|
|
if ste is not None:
|
|
t_max = max(ste.t_max, t_max)
|
|
parameter['simulation'].update({'t_max': t_max})
|
|
|
|
sim = _parse_sim(parameter['simulation'])
|
|
dist = _parse_dist(parameter['correlation_times'])
|
|
motion = _parse_motion(parameter['motion'])
|
|
mol = _parse_molecule(parameter['molecule'])
|
|
|
|
p = Parameter(sim=sim, ste=ste, spec=spec, dist=dist, motion=motion, molecule=mol)
|
|
|
|
return p
|
|
|
|
|
|
def _parse_sim(params: dict[str, Any]) -> SimParameter:
|
|
sim = SimParameter(
|
|
num_walker=params['num_walker'],
|
|
seed=params['seed'],
|
|
t_max=params['t_max']
|
|
)
|
|
return sim
|
|
|
|
|
|
def _parse_ste(params: dict[str, Any] | None) -> StimEchoParameter | None:
|
|
if params is None:
|
|
return
|
|
|
|
ste = StimEchoParameter(
|
|
t_mix=_make_times(params['t_mix']),
|
|
t_evo=_make_times(params['t_evo']),
|
|
t_echo=params.get('t_echo', 0)
|
|
)
|
|
return ste
|
|
|
|
|
|
def _parse_spectrum(params: dict[str, Any] | None) -> SpectrumParameter | None:
|
|
if params is None:
|
|
return
|
|
|
|
spec = SpectrumParameter(
|
|
num_points=params['num_points'],
|
|
dwell_time=params['dwell_time'],
|
|
t_echo=_make_times(params['t_echo']),
|
|
lb=params.get('line_broadening', 0),
|
|
t_pulse=params.get('t_pulse', 0)
|
|
)
|
|
|
|
return spec
|
|
|
|
|
|
def _parse_dist(params: dict[str, Any]) -> DistParameter:
|
|
mapping: dict = {
|
|
'DeltaDistribution': DeltaDistribution,
|
|
'LogGaussian': LogGaussianDistribution
|
|
}
|
|
p = DistParameter(
|
|
name=params['distribution'],
|
|
dist_type=mapping[params['distribution']],
|
|
variables={k: _make_times(v) for k, v in params.items() if k != 'distribution'},
|
|
)
|
|
|
|
return p
|
|
|
|
|
|
def _parse_motion(params: dict[str, Any]) -> MotionParameter:
|
|
mapping = {
|
|
'RandomJump': RandomJump,
|
|
'TetrahedralJump': TetrahedralJump,
|
|
}
|
|
|
|
p = MotionParameter(
|
|
name=params['model'],
|
|
model=mapping[params['model']],
|
|
variables={k: _make_times(v) for k, v in params.items() if k != 'model'}
|
|
)
|
|
return p
|
|
|
|
|
|
def _parse_molecule(params: dict[str, Any]) -> MoleculeParameter:
|
|
return MoleculeParameter(
|
|
delta=params['delta'],
|
|
eta=params['eta']
|
|
)
|
|
|
|
|
|
def _make_times(params: float | int | dict[str, Any]) -> np.ndarray:
|
|
times = None
|
|
|
|
if isinstance(params, (int, float, complex)):
|
|
times = np.array([params])
|
|
|
|
else:
|
|
if all(k in params for k in ('start', 'stop', 'steps')):
|
|
space_func = np.linspace
|
|
if 'is_log' in params and params['is_log']:
|
|
space_func = np.geomspace
|
|
|
|
times = space_func(start=params['start'], stop=params['stop'], num=params['steps'])
|
|
|
|
if 'list' in params:
|
|
if times is not None:
|
|
raise ValueError('list and range is given')
|
|
|
|
times = np.array(params['list'])
|
|
|
|
if times is None:
|
|
raise ValueError('No times are given')
|
|
|
|
return times
|