"""
Python api for controlling webview.
These are mostly called by webview and by the command line startup script. They
are not particularly useful from a jupyter notebook.
import bumps.names as bp
# Start the bumps server and run it in the background
await bp.start_bumps()
# Display webview in a jupyter cell
display_bumps(height=600)
# Load a model script, possibly with additional command line arguments:
path = Path("path/to/model.py")
problem = bp.load_problem(path, args=[arg1, ...])
# Use a problem defined in a separate jupyter cell
await bp.set_problem(problem, new_model=False)
"""
from functools import lru_cache
import itertools
from types import GeneratorType
from typing import (
Any,
Awaitable,
Callable,
Dict,
List,
Literal,
Mapping,
Optional,
Protocol,
Sequence,
Union,
TypedDict,
)
from datetime import datetime
import numbers
import warnings
import numpy as np
import asyncio
from pathlib import Path
import json
from copy import deepcopy
import sys
import uuid
import traceback
import time
from . import errplot
from . import options as fit_options
from .fit_thread import FitThread, EVT_FIT_COMPLETE, EVT_FIT_PROGRESS
from .fitproblem import Fitness, FitProblem, load_problem, load_pars, fitness_chisq_str
from .fitters import FitDriver, OptimizeResult, FIT_DEFAULT_ID, FIT_ACTIVE_IDS, save_fit_result
from .logger import logger
from .mapper import MPMapper
from .parameter import Parameter, Constant, Variable, unique
from .dream.stats import var_stats
from .dream.state import MCMCDraw
from .util import push_mpl_backend
from .state import (
UNDEFINED_TYPE,
FitResult,
State,
get_custom_plots_available,
serialize_problem,
deserialize_problem,
serialize_problem_bytes,
SERIALIZER_EXTENSIONS,
TopicNameType,
)
from .plots.varplot import plot_vars
from .plots.traceplot import plot_trace
from .plots.convergence_plot import convergence_plot
from .plots.custom_plot import process_custom_plot, CustomWebviewPlot
from .plots.corrplot import Corr2d
# CRUFT: os.listdrives requires python 3.12 (and only exists on windows)
try:
from os import listdrives
except ImportError:
[docs]
def listdrives() -> list[str]:
"""List Windows drive roots, e.g. ["C:\\", "D:\\"]; empty on other platforms."""
if sys.platform == "win32":
import ctypes
bitmask = ctypes.windll.kernel32.GetLogicalDrives()
return [f"{chr(65 + i)}:\\" for i in range(26) if bitmask & (1 << i)]
return []
# CRUFT: python 3.8 does not have asyncio.to_thread
try:
from asyncio import to_thread
except ImportError:
async def to_thread(func, *args, **kwargs):
"""Run a synchronous function in a separate thread."""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, func, *args, **kwargs)
REGISTRY: Dict[str, Callable] = {}
MODEL_EXT = ".json"
TRACE_MEMORY = False
# TODO: reloading the module wipes out state
# TODO: any other state that needs to be initialized?
# Initialize state with default fitter id and default options for each fitter.
state = State()
state.shared.selected_fitter = FIT_DEFAULT_ID
state.shared.fitter_settings = deepcopy(fit_options.get_fitter_defaults(active_only=False))
[docs]
def register(fn: Callable):
REGISTRY[fn.__name__] = fn
return fn
[docs]
def register_download(mimetype: str = "application/octet-stream", filename: str = "result.dat"):
"""
Decorator to register a function as a download endpoint with a specific mimetype.
"""
def decorator(fn: Callable):
# Store the mimetype as an attribute on the function
fn.mimetype = mimetype
fn.filename = filename
# Register the function in the REGISTRY
REGISTRY[fn.__name__] = fn
return fn
return decorator
[docs]
class Emitter(Protocol):
def __call__(
self,
event: str,
data: Optional[Any] = None,
to: Optional[str] = None,
room: Optional[str] = None,
skip_sid: Optional[str] = None,
namespace: Optional[str] = None,
callback: Optional[Callable] = None,
ignore_queue: bool = False,
) -> Awaitable: ...
EMITTERS: Dict[str, Emitter] = {}
[docs]
async def emit(
event: str,
data: Optional[Any] = None,
to: Optional[str] = None,
room: Optional[str] = None,
skip_sid: Optional[str] = None,
namespace: Optional[str] = None,
callback: Optional[Callable] = None,
ignore_queue: bool = False,
):
results = {}
for emitter_name, emitter_fn in EMITTERS.items():
results[emitter_name] = await emitter_fn(
event,
data=data,
to=to,
room=room,
skip_sid=skip_sid,
namespace=namespace,
callback=callback,
ignore_queue=ignore_queue,
)
return results
[docs]
@register
async def load_problem_file(
pathlist: List[str],
filename: str,
autosave_previous: bool = True,
args: List[str] = None,
):
"""
Load the problem from json or from a script file.
*pathlist* is a list of folder components and *filename* is the script file in that folder.
These are joined together as "Path(*pathlist, filename)" to build the complete path. If
path is already a Path to the file, use *load_problem_file([path.parent], path.name, ...)*
If *autosave_previous* then store the current problem state in the session file before
loading the new problem (default=True).
*args* are any additional arguments to the script file. This will be available in the script
as *sys.argv[1:]*.
"""
# print("load_problem_file", state.fitting.fit_state)
path = Path(*pathlist, filename)
logger.info(f"Loading model: {path}")
await log(f"Loading model: {path}")
problem = load_problem(path, args=args)
await set_problem(problem, Path(*pathlist), filename, new_model=True, autosave_previous=autosave_previous)
state.autosave()
[docs]
@register
async def set_serialized_problem(
serialized, new_model: bool = False, name: Optional[str] = None, method: str = "dataclass"
):
"""
Set the fit problem from a saved problem state.
*serialized* is the serialized fit problem. *method* is the method used for serialization.
If *new_model* is True, then save the model to history with tag "Loaded model". (default=False)
*name* is an optional override for the model name.
For example::
await set_serialized_problem(api.state.problem.fitProblem, method=api.state.problem.serializer)
"""
fitProblem = deserialize_problem(serialized, method=method)
await set_problem(fitProblem, new_model=new_model, name=name)
state.autosave()
[docs]
async def set_fit_result(fit_result: FitResult):
if fit_result is None:
state.reset_fitstate()
return
# Allow fit_result to be either a webview FitResult or a simple fitter OptimizeResult
# fit has more fields, and uses fit_state instead of state for the DREAM state value.
fit_result = FitResult(
method=fit_result.method,
options=fit_result.options,
convergence=fit_result.convergence,
fit_state=getattr(fit_result, "fit_state", getattr(fit_result, "state", None)),
)
# Use shared settings by default, update from any provided options
state.set_convergence(fit_result.convergence)
state.set_fit_state(fit_result.fit_state, method=fit_result.method)
state.shared.active_history = None
await set_fit_options(fitter_id=fit_result.method, options=fit_result.options)
[docs]
async def set_problem(
problem: FitProblem,
path: Optional[Path] = None,
filename: str = "",
new_model: bool = True,
name: Optional[str] = None,
autosave_previous: bool = True,
fit: Optional[OptimizeResult] = None,
):
"""
Set the fit problem.
*problem* is a fitting problem defined in a jupyter cell.
*path* and *filename* give the nominal location of the problem on disk. These are displayed
in webview and stored in the session file, but not actually used to load the problem.
*new_model* is True if the problem should be saved to the session file as "Loaded model" (default True)
*name* is an optional override for the model name.
"""
# TODO: should only save the previous if it has been modified.
# Save the old model before doing anything else
if (
autosave_previous
and state.shared.autosave_history
and state.problem is not None
and state.problem.fitProblem is not None
):
await save_to_history("autosaved before loading new model")
# if state.problem is None or state.problem.fitProblem is None:
# update = False
state.problem.fitProblem = problem
problem_path = getattr(problem, "path", None)
problem_name = getattr(problem, "name", None)
if not filename and problem_path:
filename = Path(problem_path).name
if not path and problem_path:
path = Path(problem_path).parent
name = name or problem_name or filename
state.shared.updated_model = now_string()
state.shared.updated_parameters = now_string()
state.shared.custom_plots_available = get_custom_plots_available(problem)
pathlist = list(path.parts) if path is not None else []
path_string = "(no path)" if path is None else str(path / filename)
await log(f"Model loaded: {path_string}")
state.shared.model_file = dict(filename=filename, pathlist=pathlist)
# Pick a serializer by trying dataclass and defaulting to cloudpickle
try:
_ = serialize_problem(problem, method="dataclass")
state.problem.serializer = "dataclass"
except Exception as exc:
logger.warning(f"Could not serialize problem as JSON (dataclass): {exc}, switching to cloudpickle")
state.problem.serializer = "cloudpickle"
# raise
# reset the fit state
await set_fit_result(fit)
if new_model:
state.shared.model_loaded = now_string()
if state.shared.autosave_history and state.problem is not None and state.problem.fitProblem is not None:
await save_to_history(f"Loaded model: {name}", keep=True)
await add_notification(content=path_string, title="Model loaded", timeout=2000)
[docs]
@register
async def get_history():
return state.get_history()
[docs]
@register
async def remove_history_item(name: str):
state.remove_history_item(name)
state.shared.updated_history = now_string()
[docs]
@register
async def save_to_history(
label: str,
keep: bool = False,
) -> str:
return state.save_to_history(
label,
keep=keep,
)
[docs]
@register
async def reload_history_item(name: str):
state.reload_history_item(name)
[docs]
@register
async def set_keep_history(name: str, keep: bool):
state.history.set_keep(name, keep)
state.shared.updated_history = now_string()
[docs]
@register
async def update_history_label(name: str, label: str):
state.history.update_label(name, label)
state.shared.updated_history = now_string()
[docs]
@register
async def save_problem_file(
pathlist: Optional[List[str]] = None,
filename: Optional[str] = None,
overwrite: bool = False,
):
"""
Export current problem to a file.
*pathlist* and *file*
"""
problem_state = state.problem
if problem_state is None:
logger.warning("Save failed: no problem loaded.")
return
if pathlist is None:
pathlist = problem_state.pathlist
if filename is None:
filename = problem_state.filename
if pathlist is None or filename is None:
logger.warning("no filename and path provided to save")
return {"filename": "", "check_overwrite": False}
path = Path(*pathlist)
serializer = state.problem.serializer
extension = SERIALIZER_EXTENSIONS[serializer]
# Avoid name collision with saved fit results, which may also be in .json
save_filename = f"{Path(filename).stem}-problem.{extension}"
if not overwrite and Path.exists(path / save_filename):
# confirmation needed:
return {"filename": save_filename, "check_overwrite": True}
serialized = serialize_problem_bytes(problem_state.fitProblem, method=serializer)
Path(path, save_filename).write_bytes(serialized)
await log(f"Saved: {save_filename} at path: {path}")
return {"filename": save_filename, "check_overwrite": False}
[docs]
@register
async def save_session():
state.save()
[docs]
@register
async def save_session_copy(pathlist: List[str], filename: str):
path = Path(*pathlist)
state.write_session_file(str(path / filename))
[docs]
@register_download(mimetype="application/octet-stream", filename="session.h5")
async def get_session():
return state.get_session_bytes()
[docs]
@register
async def load_session(pathlist: List[str], filename: str, read_only: bool = False):
state.setup_backing(filename, pathlist, read_only=read_only)
state.shared.updated_model = now_string()
state.shared.updated_parameters = now_string()
[docs]
@register
async def update_serialized_problem(
serialized: str,
method: str = "dataclass",
name: Optional[str] = None,
):
"""Update the current FitProblem without resetting the fit state.
Use this instead of ``set_serialized_problem`` when you want to resume
DREAM from the existing chain population rather than starting fresh.
Parameter-space changes (parameters added, removed, or renamed) are handled
automatically on the next ``start_fit_thread("dream", options, resume=True)``
call: ``DreamFit.solve`` detects the label mismatch and rebuilds the chain
state via ``_rebuild_mcmc_state`` before sampling begins.
Use ``set_serialized_problem()`` for a true cold start.
The *method* parameter accepts ``"dataclass"`` (default, safe),
``"pickle"``, ``"cloudpickle"``, or ``"dill"``. The latter three enable
arbitrary code execution and carry the same caveat as
``set_serialized_problem``.
"""
# Deserialize the new problem first so we fail fast on bad input.
fitProblem = deserialize_problem(serialized, method=method)
# Install the new problem; the existing fit_state is preserved intact.
# Any parameter-space mismatch is resolved automatically by DreamFit.solve().
state.problem.fitProblem = fitProblem
if name:
fitProblem.name = name
state.shared.updated_model = now_string()
state.shared.updated_parameters = now_string()
state.autosave()
[docs]
@register
async def set_session_output_file(filepath: Optional[Union[str, Path]] = None):
"""
Set the session output file to be used for saving results, and enable autosave.
If `filepath` is None, the session output file is cleared and autosave is disabled.
"""
if filepath is None:
await state.shared.set("session_output_file", None)
await state.shared.set("autosave_session", False)
else:
if isinstance(filepath, str):
filepath = Path(filepath)
parent_dir = filepath.parent
filename = filepath.name
if not parent_dir.exists():
raise ValueError(f"Parent directory {parent_dir} does not exist.")
if not parent_dir.is_dir():
raise ValueError(f"Path {parent_dir} is not a directory.")
if filepath.is_dir():
raise ValueError(f"Path {filepath} is a directory, not a file.")
await state.shared.set("session_output_file", dict(filename=filename, pathlist=parent_dir.parts))
await state.shared.set("autosave_session", True)
[docs]
@register
async def get_serializer():
output = {"serializer": "", "extension": ""}
problem_state = state.problem
if problem_state is not None:
serializer = problem_state.serializer
output["serializer"] = serializer
output["extension"] = SERIALIZER_EXTENSIONS[serializer]
return output
[docs]
@register
async def export_results(export_path: Union[str, List[str]] = ""):
# print("export nap"); await asyncio.sleep(0.1)
# from concurrent.futures import ThreadPoolExecutor
problem_state = state.problem
if problem_state is None or problem_state.fitProblem is None:
logger.warning("Save failed: no problem loaded.")
return
problem = deepcopy(problem_state.fitProblem)
serializer = problem_state.serializer
# TODO: if making a temporary copy of the uncertainty state is going to cause memory
# issues, we could try to copy and then fall back to just using the live object,
# or we could just always use the live object, which is unlikely to be changed before
# the export completes, anyway.
fit = deepcopy(state.fitting)
if not isinstance(export_path, list):
export_path = [export_path]
path = Path(*export_path).expanduser().absolute()
notification_id = await add_notification(content=f"<span>{str(path)}</span>", title="Export started", timeout=None)
try:
await to_thread(export_fit, path, problem, fit, serializer)
finally:
await emit("cancel_notification", notification_id)
# print("done export thread")
# Note: need naming convention for sync and async versions
[docs]
def export_fit(
path: Path | str,
problem: FitProblem,
fit: FitResult | OptimizeResult,
serializer: Optional[str] = "dataclass",
basename: Optional[str] = None,
):
# print("running export thread")
from .util import redirect_console
path = Path(path)
# Get basename for the export if not provided
if not basename:
problem_name = getattr(problem, "name", "model")
problem_path = getattr(problem, "path", f"{problem_name}.py")
basename = Path(problem_path).with_suffix("").name
# Storage directory
path.mkdir(parents=True, exist_ok=True)
output_pathstr = str(path / basename)
# Ask model to save its information
problem.save(output_pathstr)
# TODO: need the same logic in save session?
# Save a snapshot of the model that can (hopefully) be reloaded
# Complicated because it tries falling back to "cloudpickle" if it fails.
fallback = None if serializer == "dataclass" else "cloudpickle"
try:
serialized = serialize_problem_bytes(problem, serializer)
except Exception:
serialized = None
if fallback:
logger.error(f"Error serializing model with {serialized}. Trying {fallback} instead")
if serialized is None and fallback is not None:
serializer = fallback
try:
serialized = serialize_problem_bytes(problem, serializer)
except Exception as exc:
logger.error(f"Error serializing model with {serialized}: {exc}")
if serialized:
extension = SERIALIZER_EXTENSIONS[serializer]
save_filename = f"{output_pathstr}.{extension}"
try:
Path(save_filename).write_bytes(serialized)
except Exception as exc:
logger.error(f"Error writing {save_filename}: {exc}")
# Save the current state of the parameters
with redirect_console(str(path / f"{basename}.out")):
problem.show()
# TODO: add fit method and options to the par file header
# Write the pars file.
_write_pars(problem, path, f"{basename}.par")
# Write a machine-readable summary of the fit results.
try:
save_fit_result(problem, fit, path / f"{basename}-fit.json")
except Exception as exc:
logger.error(f"Error writing {basename}-fit.json: {exc}")
# Handle OptimizeResult or FitResult
fit_state = getattr(fit, "fit_state", getattr(fit, "state", None))
with push_mpl_backend("agg"):
# Produce model plots
problem.plot(figfile=output_pathstr)
# TODO: produce convergence plot and write convergence data
# Produce uncertainty plots
# TODO: Add save/show methods to the fit_state protocol
if hasattr(fit_state, "show"):
with redirect_console(str(path / f"{basename}.err")):
fit_state.show(figfile=output_pathstr)
fit_state.save(output_pathstr)
# TODO: duplicates code in fitters.DreamFit.error_plot
# TODO: refactor to separate calculation from display
# TODO: share calc_errors result with get_model_uncertainty_plot
points = errplot.error_points_from_state(state=fit_state)
res = errplot.calc_errors(problem, points)
if res is not None:
errplot.show_errors(res, save=output_pathstr)
# print("export complete")
[docs]
@register
async def save_parameters(pathlist: List[str], filename: str, overwrite: bool = False):
problem_state = state.problem
if problem_state is None:
await log("Error: Can't save parameters if no problem loaded")
return
problem = problem_state.fitProblem
path = Path(*pathlist)
if not overwrite and (path / filename).exists():
# confirmation needed:
return {"filename": filename, "check_overwrite": True}
_write_pars(problem, path, filename)
return {"filename": filename, "check_overwrite": False}
def _write_pars(problem, path: Path, filename: str):
pardata = "".join(f"{name} {value:.15g}\n" for name, value in zip(problem.labels(), problem.getp()))
with open(path / filename, "wt") as fd:
fd.write(pardata)
[docs]
@register
async def apply_parameters(pathlist: List[str], filename: str):
path = Path(*pathlist)
fullpath = path / filename
try:
# print(f"loading parameters from {fullpath}")
load_pars(state.problem.fitProblem, fullpath)
state.shared.updated_parameters = now_string()
await log(f"Applied parameters from {fullpath}")
await add_notification(
f"Applied parameters from {fullpath}",
title="Parameters applied",
timeout=2000,
)
except Exception as exc:
msg = f"error loading parameters from {fullpath}: {exc}"
logger.error(msg)
await log(msg)
await add_notification(msg, title="Error applying parameters", timeout=2000)
[docs]
@register
async def start_fit(options):
problem_state = state.problem
if problem_state is None:
await log("Error: Can't start fit if no problem loaded")
else:
fitProblem = problem_state.fitProblem
mapper = MPMapper.start_mapper(fitProblem, None, cpus=state.parallel)
monitors = []
# TODO: let FitDriver find the fitter using options["fit"]
fitclass = fit_options.lookup_fitter(options["fit"])
driver = FitDriver(
fitclass=fitclass,
mapper=mapper,
problem=fitProblem,
monitors=monitors,
**options,
)
x, fx = driver.fit()
driver.show()
[docs]
@register
async def stop_fit(wait=True):
"""
Trigger the abort fit signal to the optimizer and wait for complete (or not).
"""
if state.fit_thread is not None and state.fit_thread.is_alive():
state.fit_abort_event.set()
if wait:
await wait_for_fit_complete()
else:
state.shared.active_fit = {}
[docs]
@register
async def get_chisq(problem: Optional[FitProblem] = None, nllf=None) -> str:
if problem is None:
problem = state.problem.fitProblem
if problem is None:
return ""
return problem.chisq_str(nllf=nllf) # Default is norm=True and compact=True
[docs]
def get_running_loop():
try:
return asyncio.get_running_loop()
except RuntimeError:
return None
[docs]
@register
async def shake_parameters():
fitProblem = state.problem.fitProblem if state.problem is not None else None
if fitProblem is not None:
# TODO: capture and report seed?
fitProblem.randomize()
state.shared.updated_parameters = now_string()
await log(f"Randomize parameters")
await add_notification(
f"Randomize parameters",
title="Parameters applied",
timeout=2000,
)
[docs]
@register
async def start_fit_thread(
fitter_id: Optional[str] = None, options: Optional[Dict[str, Any]] = None, resume: bool = False
):
fitProblem = state.problem.fitProblem if state.problem is not None else None
if fitProblem is None:
await log("Error: Can't start fit if no problem loaded")
return
state.calling_loop = get_running_loop()
if state.fit_thread is not None:
# warn that fit is alread running...
logger.warning("fit already running...")
await log("Can't start fit, a fit is already running...")
return
# TODO: better access to model parameters
num_params = len(fitProblem.getp())
if num_params == 0:
raise ValueError("Problem has no fittable parameters")
# Check the options. Pass the fitter_id so that we know which options are available.
if options is None:
options = {}
# Allow fit=fitter_id or method=fitter_id in the options dictionary.
# Using None as the default so that fitoptions.check_options an fill
# in any a value if the fit id was not specified.
id_from_options = options.pop("fit", options.pop("method", None))
if not fitter_id:
fitter_id = id_from_options
options, errors = fit_options.check_options(options, fitter_id=fitter_id)
for msg in errors:
logger.warning(msg)
await log(msg)
# Start a new thread worker and give fit problem to the worker.
# Clear abort and uncertainty state
# state.abort = False
# state.fitting.uncertainty_state = None
state.fit_abort_event.clear()
# TODO: remove this re-creation of the Event object when minimum python is >= 3.10
state.fit_complete_event = asyncio.Event()
state.fit_complete_event.clear()
# Use shared settings by default, update from any provided options
shared_settings = state.shared.fitter_settings
full_options = shared_settings[fitter_id]["settings"].copy()
if options:
full_options.update(options)
# TODO: model.py may have changed; check that the list of parameters is the same
# TODO: maybe prefer problem saved in store on resume
# print(f"start fit thread {resume} {state.fitting.fit_state}")
fitclass = fit_options.lookup_fitter(fitter_id)
max_steps = fitclass.max_steps(fitProblem, full_options)
if resume and state.fitting.method != fitter_id:
msg = f"Can't resume {fitter_id} from state saved by {state.fitting.method}"
logger.warning(msg)
await log(msg)
resume = False
state.reset_fitstate(copy=resume)
state.fitting.method = fitter_id
state.fitting.options = full_options
state.shared.active_fit = to_json_compatible_dict(
dict(
fitter_id=fitter_id,
options=full_options,
num_steps=max_steps,
# TODO: step should be length fitting.convergence on resume
step=0,
chisq="",
value=0,
)
)
fit_thread = FitThread(
fit_abort_event=state.fit_abort_event,
fitclass=fitclass,
problem=fitProblem,
mapper=state.mapper,
options=full_options,
parallel=state.parallel,
# session_id=session_id,
# Number of seconds between updates to the GUI, or 0 for no updates
convergence_update=5,
uncertainty_update=state.shared.autosave_session_interval,
console_update=state.console_update_interval,
fit_state=state.fitting.fit_state,
convergence=state.fitting.convergence,
)
await log(
json.dumps(to_json_compatible_dict(options), indent=2),
title=f"Starting fitter {fitter_id}",
)
state.autosave()
fit_thread.start()
state.fit_thread = fit_thread
[docs]
@register
async def set_fit_options(fitter_id: str, options: Dict[str, Any]):
current_options = state.shared.fitter_settings[fitter_id]["settings"]
current_options.update(options)
# TODO: do we need to update state.fitting.options as well?
# state.fitting.options = current_options.copy()
# items in state.shared are not deeply reactive, so we have to explicitly notify:
state.shared.notify("fitter_settings")
[docs]
async def wait_for_fit_complete():
if state.fit_thread is not None:
await state.fit_complete_event.wait()
async def _fit_progress_handler(event: Dict):
# print("inside _fit_progress_handler", event)
# session_id = event["session_id"]
if TRACE_MEMORY:
import tracemalloc
if tracemalloc.is_tracing():
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics("lineno")
print("memory use:")
for stat in top_stats[:15]:
print(stat)
problem_state = state.problem
fitProblem = problem_state.fitProblem if problem_state is not None else None
if fitProblem is None:
raise ValueError("should never happen: fit progress reported for session in which fitProblem is undefined")
message = event.get("message", None)
# print("_fit_progress_handler", message)
if message == "complete" or message == "improvement":
fitProblem.setp(event["point"])
fitProblem.model_update()
state.shared.updated_parameters = now_string()
if message == "complete":
state.shared.active_fit = {}
elif message == "convergence_update":
state.set_convergence(event["convergence"])
elif message == "progress":
active_fit = state.shared.active_fit
active_fit.update({"step": event["step"], "chisq": event["chisq"]})
state.shared.active_fit = active_fit
elif message == "uncertainty_update" or message == "uncertainty_final":
state.set_fit_state(event["fit_state"], event["method"])
if message != "uncertainty_final":
# don't save state for uncertainty_final- the fit_complete handler will do that.
state.autosave()
async def _fit_complete_handler(event: Dict[str, Any]):
message = event.get("message", None)
# print("inside _fit_complete_handler", message)
try:
if state.fit_thread is not None:
# print("joining fit thread")
state.fit_thread.join(1) # 1 second timeout on join
if state.fit_thread.is_alive():
# TODO: what can we do to force quit the thread?
await log("Fit thread failed to complete")
# return ?
# print("...joined")
if message == "error":
await log(
event["traceback"],
title=f"Fit failed with error: {event['error_string']}",
)
logger.warning(f"Fit failed with error: {event['error_string']}\n{event['traceback']}")
return
# print(event['info']) # Needed if we are dumping fit outputs to the terminal
problem: FitProblem = event["problem"]
chisq = nice(problem.chisq(nllf=event["value"]))
problem.setp(event["point"]) # setp calls model_update
state.problem.fitProblem = problem
state.set_fit_state(event["fit_state"], event["fitter_id"])
if state.shared.autosave_history:
item_timestamp = await save_to_history(
f"Fit complete: {event['fitter_id']}",
)
state.shared.active_history = item_timestamp
state.autosave()
state.shared.updated_parameters = now_string()
await log(event["info"], title=f"Done with chisq {chisq}")
logger.info(f"Fit done with chisq {chisq}")
finally:
# Signal that the fit is complete and all results are saved.
# Be sure to clear the fit thread and active fit before so that those
# awaiting the fit complete event can resume and start a new fit.
state.fit_thread = None
state.shared.active_fit = {}
# Signal to those waiting that the fit is complete.
state.fit_complete_event.set()
# print("shutdown", state.shutdown_on_fit_complete)
if state.shutdown_on_fit_complete:
await shutdown()
# print("shutdown complete")
[docs]
def call_async(async_fn, *args, **kw):
"""
Call an async function inside the active loop, reporting any exceptions
on the logger.
"""
# print(f"call async {async_fn}") # (*({args}), **kw({kw}))")
loop = getattr(state, "calling_loop", None)
if loop is not None:
async def trap_exceptions():
# print("inside exception trap")
try:
await async_fn(*args, **kw)
except Exception as exc:
logger.exception(exc)
task = asyncio.run_coroutine_threadsafe(trap_exceptions(), loop)
# task.result(120)
[docs]
def fit_progress_handler(event: Dict):
call_async(_fit_progress_handler, event)
[docs]
def fit_complete_handler(event: Dict):
call_async(_fit_complete_handler, event)
# Run from the fit thread by blink. The handlers echo the message to asyncio
# handlers for communication with the GUI and in the case of FIT_COMPLETE, for
# saving the results.
EVT_FIT_PROGRESS.connect(fit_progress_handler, weak=True)
EVT_FIT_COMPLETE.connect(fit_complete_handler, weak=True)
[docs]
async def log(message: str, title: Optional[str] = None):
topic = "log"
contents = {
"message": {"message": message, "title": title},
"timestamp": now_string(),
}
# session = get_session(session_id)
state.topics[topic].append(contents)
# if session_id == app["active_session"]:
await emit(topic, contents)
[docs]
@register
async def get_data_plot(model_indices: Optional[List[int]] = None):
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = deepcopy(state.problem.fitProblem)
# Suppress all mpld3 warnings
# warnings.filterwarnings("ignore", module="mpld3")
# TODO: revise get_data_plot interface to take only a single index
# The current interface to get_data_plot takes a list of model indices
# but it only returns a single figure. If called with a list of models
# they will overwrite each other in one figure.
if model_indices is None or len(model_indices) != 1:
raise RuntimeError("can only do one model at a time")
index = model_indices[0]
# Overall chisq
overall_chisq_str = fitProblem.chisq_str()
with fitProblem.push_model(index) as model:
# # "per model" chisq
if fitProblem.num_models > 1:
chisq_str = fitness_chisq_str(model)
text = f"χ² = {chisq_str}; overall {overall_chisq_str}"
title = f"Model {index+1}: {model.name}"
else:
text = f"χ² = {overall_chisq_str}"
title = f"{model.name}"
if hasattr(model, "plotly"):
return _get_data_plot_plotly(model, title=title, chisq=text)
elif hasattr(model, "plot"):
return _get_data_plot_mpl(model, title=title, chisq=text)
else:
# Use plotly to show chisq
return _get_data_plot_plotly(model, title=title, chisq=text)
def _get_data_plot_plotly(model, title, chisq):
if hasattr(model, "plotly"):
fig = model.plotly()
else:
import plotly.graph_objects as go
fig = go.Figure()
fig.update_layout(
xaxis_visible=False,
yaxis_visible=False,
# plot_bgcolor='rgba(0,0,0,0)',
# paper_bgcolor='rgba(0,0,0,0)',
)
# TODO: text offset of (x=0.5em, y=0.5ex)
text_offset = 0.01 # portion of graph axis length
font = dict(size=22)
# fig.add_annotation(
# x=text_offset, y=1+text_offset,
# xanchor="left", yanchor="bottom",
# xref="paper", yref="paper",
# text=title,
# showarrow=False,
# font=font,
# )
fig.add_annotation(
x=1 - text_offset,
y=1 + text_offset,
xanchor="right",
yanchor="bottom",
xref="paper",
yref="paper",
text=chisq,
showarrow=False,
font=font,
)
dfig = fig.to_dict()
return {"fig_type": "plotly", "plotdata": to_json_compatible_dict(dfig)}
# Make mpld3 figure controls available for monkey-patching.
# Figure size is a hack. The correct choice will probably depend on browser,
# screen size, dpi, available fonts, etc.
# TODO: have webview send the desired figure size
MPLD3_BACKEND = "agg"
MPLD3_FIG_SIZE = (10, 8)
MPLD3_STYLE = {
"figure.dpi": 72,
"figure.subplot.left": 0.10,
"figure.subplot.right": 0.99,
"figure.subplot.bottom": 0.05,
"figure.subplot.top": 0.95,
"axes.xmargin": 0.05,
"axes.ymargin": 0.05,
"figure.constrained_layout.h_pad": 0.0,
"figure.constrained_layout.w_pad": 0.0,
"figure.constrained_layout.hspace": 0.0,
"figure.constrained_layout.wspace": 0.0,
"font.size": 16,
}
def _get_data_plot_mpl(model, title, chisq):
import matplotlib.pyplot as plt
import mpld3
from mpld3 import plugins
start_time = time.time()
logger.info(f"queueing new data plot... {start_time}")
text_offset = 0.01 # portion of graph axis length
# Note: rc_context() says it won't modify the backend, so we still need push_mpl_backend().
with push_mpl_backend(MPLD3_BACKEND), plt.rc_context(MPLD3_STYLE):
fig = plt.figure(figsize=MPLD3_FIG_SIZE)
model.plot()
# TODO: can't adjust margins correctly unless we know the figure size
# h, w = fig.get_size_inches()
# h_ex = h*72 / 16 # (h in * 72 pt/in) / (16 pt/ex) = height in ex
h_ex = 30 # assume we are 50 lines tall, so that 2/30 ~ 0.08
text_offset = 0.5 / h_ex # 1/2 ex above and below the text
top = 1 - 2 / h_ex # leave 2 ex at the top of the figure
plt.subplots_adjust(top=top)
# # transFigure doesn't seem to work in mpld3
# transform = fig.transFigure
# x, y = text_offset, 1 - text_offset
# ha, va = "left", "top"
# fig.text(x, y, title, transform=transform, va=va, ha=ha)
# x, y = 1 - text_offset, 1 - text_offset
# ha, va = "right", "top"
# fig.text(x, y, chisq, transform=transform, va=va, ha=ha)
ax = fig.axes[0]
fontsize = 22
transform = ax.transAxes
# Don't need the title since it is in the models dropdown
# x, y = text_offset, 1 + text_offset
# ha, va = "left", "bottom"
# ax.text(x, y, title, transform=transform, va=va, ha=ha, fontsize=fontsize)
x, y = 1 - text_offset, 1 + text_offset
ha, va = "right", "bottom"
ax.text(x, y, chisq, transform=transform, va=va, ha=ha, fontsize=fontsize)
# Add plugins to the figure for zoom, etc.
# Note: these are already present in the dict, just not working on the backend?
plugins.clear(fig)
plugins.connect(
fig,
plugins.BoxZoom(enabled=True),
plugins.Reset(),
plugins.MousePosition(fontsize=14),
)
dfig = mpld3.fig_to_dict(fig)
# import pprint; pprint.pprint(dfig); import sys; sys.exit()
plt.close(fig)
end_time = time.time()
logger.info(f"time to draw data plot: {end_time - start_time}")
return {"fig_type": "mpld3", "plotdata": to_json_compatible_dict(dfig)}
[docs]
@register
async def get_model_names():
problem = state.problem.fitProblem
if problem is None:
return None
return [p.name if p.name is not None else f"model_{i}" for (i, p) in enumerate(problem.models)]
[docs]
@register
async def get_model():
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = state.problem.fitProblem
serialized = serialize_problem(fitProblem, "dataclass") if state.problem.serializer == "dataclass" else "null"
return serialized
[docs]
class WebviewPlotFunction(Protocol):
def __call__(
self,
model: Fitness,
problem: FitProblem,
state: Optional[MCMCDraw] = None,
n_samples: Optional[int] = None,
) -> dict: ...
# custom plots are an opt-in feature for models
# they are defined in the model file as a dictionary of functions
# with a "change_with" key that specifies whether the plot should
# change with the uncertainty state or with the parameters
[docs]
@register
async def get_custom_plot_info():
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = state.problem.fitProblem
output: List[dict] = []
for model_index, model in enumerate(fitProblem.models):
model_webview_plots = getattr(model, "webview_plots", {})
for title, plotinfo in model_webview_plots.items():
output.append(
{
"model_index": model_index,
"change_with": plotinfo["change_with"],
"title": title,
}
)
return output
[docs]
async def create_custom_plot(model_index: int, plot_title: str, n_samples: int = 1) -> CustomWebviewPlot:
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = deepcopy(state.problem.fitProblem)
fit_state = state.fitting.fit_state
# update model
model = list(fitProblem.models)[model_index]
webview_plots = getattr(model, "webview_plots", {})
plot_info = webview_plots.get(plot_title, {})
plot_function: WebviewPlotFunction = webview_plots.get(plot_title, {}).get("func", None)
if plot_function is not None:
try:
model.update()
model.nllf()
if plot_info.get("change_with", None) == "uncertainty":
plot_item: CustomWebviewPlot = await to_thread(plot_function, model, fitProblem, fit_state, n_samples)
else:
plot_item: CustomWebviewPlot = await to_thread(plot_function, model, fitProblem)
except Exception:
plot_item = CustomWebviewPlot(fig_type="error", plotdata=traceback.format_exc(), exportdata=None)
return process_custom_plot(plot_item)
return {}
[docs]
@register
async def get_custom_plot(model_index: int, plot_title: str, n_samples: int = 1):
output = CustomWebviewPlot(figtype="error", plotdata="no plot")
if model_index is not None:
figdict = await create_custom_plot(model_index=model_index, plot_title=plot_title, n_samples=n_samples)
output = to_json_compatible_dict(figdict)
return output
[docs]
@register
async def get_convergence_plot(
cutoff: float = 0.25, portion: Optional[float] = None, max_points: Optional[int] = 10000
):
"""
Get the convergence plot for the current fit state.
If the fit state is not available, return None.
If the convergence is not available, return None.
:param cutoff: The cutoff value for the convergence plot
(fraction of points below this value are not shown)
:param max_points: The maximum number of points to plot
(thinning applied if too many points)
:return: A JSON-serializable dictionary containing the convergence plot data."""
if state.problem is None or state.problem.fitProblem is None:
return None
dof = state.problem.fitProblem.dof
convergence = state.fitting.convergence
if convergence is not None:
fit_state = state.fitting.fit_state
generation = len(convergence)
if fit_state is not None and hasattr(fit_state, "trim_index"):
# If the trim index is available, we can show it on the plot:
trim_index = fit_state.trim_index(generation=generation, portion=portion)
burn_index = fit_state.trim_index(generation=generation, portion=1.0)
stored_portion = getattr(fit_state, "portion", None)
else:
trim_index = None
burn_index = None
stored_portion = None
plotdata = convergence_plot(
convergence, dof, cutoff=cutoff, trim_index=trim_index, burn_index=burn_index, max_points=max_points
)
output = {
"plotdata": plotdata,
"portion": stored_portion,
}
return to_json_compatible_dict(output)
else:
return None
[docs]
@register
async def set_trim_portion(portion: float):
"""
Set the trim portion for the current fit state.
This will update the trim index and burn index in the fit state.
"""
fit_state = state.fitting.fit_state
if fit_state is not None and hasattr(fit_state, "portion"):
if not (0.0 <= portion <= 1.0):
raise ValueError("Trim portion must be between 0.0 and 1.0")
fit_state.portion = portion
state.shared.updated_convergence = now_string()
state.shared.updated_uncertainty = now_string()
await add_notification(
f"Set trim portion to {portion}",
title="Trim portion set",
timeout=2000,
)
@lru_cache(maxsize=30)
def _get_correlation_plot(
sort: bool = True,
max_rows: int = 8,
nbins: int = 50,
vars=None,
timestamp: str = "",
):
fit_state = state.fitting.fit_state
if hasattr(fit_state, "draw"):
start_time = time.time()
logger.info(f"queueing new correlation plot... {start_time}")
draw = fit_state.draw(vars=vars)
c = Corr2d(draw.points.T, bins=nbins, labels=draw.labels)
fig = c.plot(sort=sort, max_rows=max_rows)
logger.info(f"time to render but not serialize... {time.time() - start_time}")
serialized = to_json_compatible_dict(fig.to_dict())
end_time = time.time()
logger.info(f"time to draw correlation plot: {end_time - start_time}")
return serialized
else:
return None
[docs]
@register
async def get_correlation_plot(
sort: bool = True,
max_rows: int = 8,
nbins: int = 50,
vars=None,
timestamp: str = "",
):
# need vars to be immutable (hashable) for caching based on arguments:
vars = tuple(vars) if vars is not None else None
result = await to_thread(
_get_correlation_plot,
sort=sort,
max_rows=max_rows,
nbins=nbins,
vars=vars,
timestamp=timestamp,
)
return result
@lru_cache(maxsize=30)
def _get_uncertainty_plot(timestamp: str = "", cbar_colors: int = 8):
fit_state = state.fitting.fit_state
if hasattr(fit_state, "draw"):
start_time = time.time()
logger.info(f"queueing new uncertainty plot... {start_time}")
draw = fit_state.draw()
nbins = max(min(draw.points.shape[0] // 20000, 100), 30)
stats = var_stats(draw)
fig = plot_vars(draw, stats, nbins=nbins, cbar_colors=cbar_colors)
logger.info(f"time to draw uncertainty plot: {time.time() - start_time}")
return to_json_compatible_dict(fig)
else:
return None
[docs]
@register
async def get_uncertainty_plot(timestamp: str = ""):
result = await to_thread(_get_uncertainty_plot, timestamp=timestamp, cbar_colors=8)
return result
[docs]
@register
async def get_model_uncertainty_plot():
import mpld3
import matplotlib.pyplot as plt
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = state.problem.fitProblem
fit_state = state.fitting.fit_state
if not hasattr(fit_state, "draw"):
return
start_time = time.time()
logger.info(f"queueing new model uncertainty plot... {start_time}")
points = errplot.error_points_from_state(fit_state)
errs = errplot.calc_errors(fitProblem, points)
logger.info(f"errors calculated: {time.time() - start_time}")
with push_mpl_backend("agg"):
fig = plt.figure()
errplot.show_errors(errs, fig=fig)
logger.info(f"time to render but not serialize... {time.time() - start_time}")
fig.canvas.draw()
dfig = mpld3.fig_to_dict(fig)
plt.close(fig)
end_time = time.time()
logger.info(f"time to draw model uncertainty plot: {end_time - start_time}")
return dfig
[docs]
@register
async def get_parameter_labels():
# Required to support get_parameter_trace_plot because ordering must be preserved.
# There is no way to know whether a disambiguated name occurred first or second from
# get_parameters. Uses fitProblem because uncertainty state might not exist and
# list of parameters should be updated on model_loaded.
# Should probably be able to call these parameters by ID.
if state.problem is None or state.problem.fitProblem is None:
return None
return to_json_compatible_dict(state.problem.fitProblem.labels())
[docs]
@register
async def get_parameter_trace_plot(var: int):
fit_state = state.fitting.fit_state
# TODO: parallel tempering has a different trace plot
if isinstance(fit_state, MCMCDraw):
import time
start_time = time.time()
logger.info(f"queueing new parameter_trace plot... {start_time}")
# Portion defaults to burn point
# Set outliers to true to include bad chains
genid, chains = fit_state.traces(portion=None, thin=1, outliers=False)
label = fit_state.labels[var]
fig = plot_trace(
genid,
chains[:, :, var],
label=label,
alpha=0.4,
)
logger.info(f"time to render but not serialize... {time.time() - start_time}")
dfig = fig.to_dict()
end_time = time.time()
logger.info(f"time to draw parameter_trace plot: {end_time - start_time}")
return to_json_compatible_dict(dfig)
else:
return None
[docs]
@register
async def get_parameters(only_fittable: bool = False):
if state.problem is None or state.problem.fitProblem is None:
return []
fitProblem = state.problem.fitProblem
freevars = fitProblem.freevars
all_parameters = fitProblem.model_parameters()
if only_fittable:
parameter_infos = params_to_list(unique(all_parameters), freevars=freevars)
# only include params with priors:
parameter_infos = [pi for pi in parameter_infos if pi["fittable"] and not pi["fixed"]]
else:
parameter_infos = params_to_list(all_parameters, freevars=freevars)
return to_json_compatible_dict(parameter_infos)
[docs]
@register
async def set_parameter(
parameter_id: str,
property: Literal["value01", "value", "min", "max"],
value: Union[float, str, bool],
):
if state.problem is None or state.problem.fitProblem is None:
return None
fitProblem = state.problem.fitProblem
parameter = fitProblem._parameters_by_id.get(parameter_id, None)
if parameter is None:
warnings.warn(f"Attempting to update parameter that doesn't exist: {parameter_id}")
return
# TODO: should never happen
if parameter.prior is None:
warnings.warn(f"Attempting to set prior properties on parameter without priors: {parameter}")
return
if property == "value01":
new_value = parameter.prior.put01(value)
nice_new_value = nice(new_value, digits=VALUE_PRECISION)
parameter.clip_set(nice_new_value)
elif property == "value":
new_value = float(value)
nice_new_value = nice(new_value, digits=VALUE_PRECISION)
parameter.clip_set(nice_new_value)
elif property == "min":
lo = float(value)
hi = parameter.prior.limits[1]
parameter.range(lo, hi)
parameter.reset_prior()
elif property == "max":
lo = parameter.prior.limits[0]
hi = float(value)
parameter.range(lo, hi)
parameter.reset_prior()
elif property == "fixed":
if parameter.fittable:
parameter.fixed = bool(value)
fitProblem.model_reset()
# logger.info(f"setting parameter: {parameter}.fixed to {value}")
# model has been changed: setp and getp will return different values!
state.shared.updated_model = now_string()
# Reset the fitting state (uncertainty and population), no longer valid
state.reset_fitstate()
fitProblem.model_update()
state.shared.updated_parameters = now_string()
return
[docs]
def now_string():
return f"{datetime.now().timestamp():.6f}"
[docs]
@register
async def publish(topic: str, message: Any = None):
timestamp_str = f"{datetime.now().timestamp():.6f}"
contents = {"message": message, "timestamp": timestamp_str}
# session = get_session(session_id)
state.topics[topic].append(contents)
# if session_id == app["active_session"]:
await emit(topic, contents)
# logger.info(f"emitted: {topic} :: {contents}")
[docs]
@register
async def get_shared_setting(setting: str):
value = await state.shared.get(setting)
return to_json_compatible_dict(value)
[docs]
@register
async def set_shared_setting(setting: str, value: Any):
await state.shared.set(setting, value)
[docs]
async def notify_shared_setting(setting: str, value: Any):
await emit(setting, to_json_compatible_dict(value))
state.shared._notification_callbacks["emit"] = notify_shared_setting
[docs]
@register
async def get_topic_messages(topic: Optional[TopicNameType] = None, max_num=None) -> List[Dict]:
# this is a GET request in disguise -
# emitter must handle the response in a callback,
# as no separate response event is emitted.
if topic is None:
return []
topics = state.topics
q = topics.get(topic, None)
if q is None:
raise ValueError(f"Topic: {topic} not defined")
elif max_num is None:
return list(q)
else:
q_length = len(q)
start = max(q_length - max_num, 0)
return list(itertools.islice(q, start, q_length))
DIRLISTING_TIMEOUT = 10.0 # seconds before an unreachable path returns an error instead of hanging
def _get_dirlisting_sync(pathlist: Optional[List[str]] = None):
"""Walk a directory synchronously; run via asyncio.to_thread (filesystem I/O may be slow on network drives)."""
subfolders = []
files = []
path = Path(state.base_path) if (pathlist is None or len(pathlist) == 0) else Path(*pathlist)
missing = None
if not path.exists():
missing = path
path = Path.cwd()
abs_path = path.absolute()
for p in abs_path.iterdir():
try:
stat = p.stat()
except OSError:
# entry vanished between iterdir() and stat(), or is a broken symlink
continue
mtime = stat.st_mtime
fileinfo = {"name": p.name, "modified": mtime}
if p.is_dir():
# NOTE: not counting directory contents — a glob per subfolder costs one
# round-trip each on network filesystems, and the count is not displayed.
fileinfo["size"] = 0
subfolders.append(fileinfo)
else:
# files.append(p.resolve().name)
fileinfo["size"] = stat.st_size
files.append(fileinfo)
# for Windows: list drives as well
drives = listdrives()
return dict(drives=drives, pathlist=abs_path.parts, subfolders=subfolders, files=files), missing
[docs]
@register
async def get_dirlisting(pathlist: Optional[List[str]] = None):
# GET request
# TODO: use psutil to get disk listing as well?
try:
result, missing = await asyncio.wait_for(
asyncio.to_thread(_get_dirlisting_sync, pathlist),
timeout=DIRLISTING_TIMEOUT,
)
except asyncio.TimeoutError:
path_str = str(Path(*pathlist)) if pathlist else str(state.base_path)
return {"error": f"Timed out reading directory: {path_str} (network drive may be unreachable)"}
except OSError as exc:
return {"error": f"Cannot read directory: {exc}"}
if missing is not None:
await add_notification(
f"Path does not exist: {missing}, falling back to current working directory",
title="Error",
timeout=2000,
)
return result
[docs]
@register
async def get_fitter_defaults():
# TODO: Need a better way to include experimental fitters in the web ui.
# This hack allows gui reset of fit options when --fit=pt is on the command line
# by allowing every experimental and deprecated fitter to appear in the fit selector.
# That all of them appear is okay since this only affects users who already know that
# there are experimental fitters available. It will probably break when we refactor
# the command line processor, or when we start webview from a jupyter notebook.
active_only = state.shared.selected_fitter in FIT_ACTIVE_IDS
return fit_options.get_fitter_defaults(active_only=active_only)
[docs]
@register
async def get_fit_fields():
return fit_options.get_fit_fields()
[docs]
@register
async def shutdown():
logger.info("killing...")
await stop_fit()
state.autosave()
if state.mapper is not None:
state.mapper.stop_mapper()
state.mapper = None
# print("gather _shutdown()")
# TODO: why gather here rather than await?
asyncio.gather(_shutdown(), return_exceptions=True)
[docs]
async def add_notification(content: str, title: str = "Notification", timeout: Optional[int] = None):
id = None
if timeout is None:
id = str(uuid.uuid4())
await emit("add_notification", {"title": title, "content": content, "id": id})
else:
await emit("add_notification", {"title": title, "content": content, "timeout": timeout})
return id
async def _shutdown():
# print("raising SystemExit")
# raise SystemExit(0)
...
VALUE_PRECISION = 6
VALUE_FORMAT = "{{:.{:d}g}}".format(VALUE_PRECISION)
[docs]
def nice(v, digits=4):
"""Fix v to a value with a given number of digits of precision"""
from math import log10, floor
v = float(v)
if v == 0.0 or not np.isfinite(v):
return v
else:
sign = v / abs(v)
place = floor(log10(abs(v)))
scale = 10 ** (place - (digits - 1))
return sign * floor(abs(v) / scale + 0.5) * scale
JSON_TYPE = Union[str, float, bool, None, Sequence["JSON_TYPE"], Mapping[str, "JSON_TYPE"]]
"""Type union of a JSON compatible python structure"""
[docs]
def to_json_compatible_dict(obj) -> JSON_TYPE:
if isinstance(obj, (list, tuple)):
return type(obj)(to_json_compatible_dict(v) for v in obj)
elif isinstance(obj, GeneratorType):
return list(to_json_compatible_dict(v) for v in obj)
elif isinstance(obj, dict):
return type(obj)((to_json_compatible_dict(k), to_json_compatible_dict(v)) for k, v in obj.items())
elif isinstance(obj, np.ndarray) and obj.dtype.kind in ["f", "i"]:
return obj.tolist()
elif isinstance(obj, np.ndarray) and obj.dtype.kind == "O":
return to_json_compatible_dict(obj.tolist())
elif isinstance(obj, bool) or isinstance(obj, str) or obj is None:
return obj
elif isinstance(obj, numbers.Number):
return str(obj) if not np.isfinite(obj) else float(obj)
elif isinstance(obj, UNDEFINED_TYPE):
return None
else:
raise ValueError("obj %s is not serializable" % str(obj))
[docs]
class ParamInfo(TypedDict, total=False):
id: str
name: str
paths: List[str]
slot_repr: str
value_str: str
fittable: bool
fixed: bool
writable: bool
value01: float
min_str: str
max_str: str
tags: List[str]
[docs]
def params_to_list(params, lookup=None, path="", freevars=None) -> List[ParamInfo]:
lookup: Dict[str, ParamInfo] = {} if lookup is None else lookup
# print(f"{type(params)=} {params=}\n {lookup=}\n {pathlist=}")
if params is None:
pass
elif isinstance(params, dict):
for k in sorted(params.keys()):
item = f"{path}{'.' if path else ''}{k}"
params_to_list(params[k], lookup=lookup, path=item, freevars=freevars)
elif isinstance(params, (tuple, list, np.ndarray)):
# print("list branch")
for i, v in enumerate(params):
params_to_list(v, lookup=lookup, path=f"{path}[{i:d}]", freevars=freevars)
elif isinstance(params, (Parameter, Constant)):
# path = ".".join(pathlist)
# print(f"failing with {params} {lookup}")
pid = params.id
has_slot = hasattr(params, "slot")
existing = lookup.get(pid, None)
if existing is not None:
existing["paths"].append(path)
elif freevars is not None and freevars.isfree(params):
# Don't include free parameters from the model in the list
# of available parameters.
pass
else:
value_str = VALUE_FORMAT.format(nice(params.value))
writable = has_slot and isinstance(params.slot, (float, Variable, Parameter))
new_item: ParamInfo = {
"id": pid,
"name": str(params.name),
"paths": [path],
"slot_repr": str(params.slot) if has_slot else "",
"tags": getattr(params, "tags", []),
"writable": writable,
"value_str": value_str,
"fittable": params.fittable,
"fixed": params.fixed,
}
if hasattr(params, "prior"):
lo, hi = params.prior.limits
new_item["value01"] = params.prior.get01(float(params.value))
new_item["min_str"] = VALUE_FORMAT.format(nice(lo))
new_item["max_str"] = VALUE_FORMAT.format(nice(hi))
lookup[pid] = new_item
# Check for any additional parameters referenced by the slot
if has_slot:
subparams = params.slot.parameters()
if subparams:
if len(subparams) > 1:
params_to_list(subparams, lookup=lookup, path=f"{params}", freevars=freevars)
else:
params_to_list(subparams[0], lookup=lookup, path="", freevars=freevars)
elif callable(getattr(params, "parameters", None)):
# handle Expression, etc.
subparams = params.parameters()
subparams = subparams[0] if subparams and len(subparams) == 1 else subparams
params_to_list(subparams, lookup=lookup, path=f"{path}={params}", freevars=freevars)
return list(lookup.values())