python/rwsims/parsing.py

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['t_echo']
    )
    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(
        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(
        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