kalpa/kalpaa/common/model_params.py
Deepak Mallubhotla a731a81c6a
Some checks failed
gitea-physics/kalpa/pipeline/head There was a failure building this commit
feat: many disparate updates for modernising
2025-02-21 15:58:12 -06:00

78 lines
1.6 KiB
Python

import logging
from tantri.dipoles.types import Orientation
import kalpaa.config
from pdme.model import (
LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel,
LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel,
LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel,
)
_logger = logging.getLogger(__name__)
def long_orientation_name(orientation: Orientation) -> str:
return {
Orientation.RANDOM: "free",
Orientation.XY: "fixedxy",
Orientation.Z: "fixedz",
}[orientation]
def _fixed_z_model_func(
xmin,
xmax,
ymin,
ymax,
zmin,
zmax,
wexp_min,
wexp_max,
pfixed,
n_max,
prob_occupancy,
):
return LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
xmin,
xmax,
ymin,
ymax,
zmin,
zmax,
wexp_min,
wexp_max,
pfixed,
0,
0,
n_max,
prob_occupancy,
)
def get_model(params: kalpaa.config.ReducedModelParams):
model_funcs = {
Orientation.Z: _fixed_z_model_func,
Orientation.RANDOM: LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel,
Orientation.XY: LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel,
}
_logger.info(f"Got params that look like {params=}")
_logger.info(f"Got params that look like {params=}")
model = model_funcs[params.orientation](
params.x_min,
params.x_max,
params.y_min,
params.y_max,
params.z_min,
params.z_max,
params.w_log_min,
params.w_log_max,
10**params.log_magnitude,
params.count,
0.99999999,
)
return (
f"geom_{model.xmin}_{model.xmax}_{model.ymin}_{model.ymax}_{model.zmin}_{model.zmax}-magnitude_{params.log_magnitude}-orientation_{long_orientation_name(params.orientation)}-dipole_count_{params.count}",
model,
)