Use new Generic bumps FitProblem (#258)

* use generic FitProblem with FitnessType=Experiment to generate schema

* add more items to Probe types, for schema generation

* use generic FitProblem with Experiment model-type specifier for typing outputs of load_model

* remove typing shim for Refl1D FitProblem, replaced by generic FitProblem[Experiment] where models are typed as Experiment

* bumps_interface.fitproblem is removed, don't include it in documentation API

doc/genmods.py

@@ -41,7 +41,6 @@
     ("deprecated.freeform", "Freeform - Parametric B-Spline"),
     ("backends", "Backends for low level reflectivity calculations"),
     ("bumps_interface.fitplugin", "Bumps plugin definition for reflectivity models"),
-    ("bumps_interface.fitproblem", "Bumps FitProblem definition for reflectivity models"),
     ("dist", "Non-uniform samples"),
     ("experiment", "Reflectivity fitness function"),
     ("probe", "Instrument probe"),

extra/gen_schema_v2.py

@@ -6,29 +6,33 @@
 
 from bumps.util import NumpyArray
 
+from refl1d.experiment import Experiment
 from refl1d.names import FitProblem
 
 NDArray = NumpyArray
 
-NDArray = NumpyArray
-
 
 class BumpsGenerateJsonSchema(GenerateJsonSchema):
     def dataclass_schema(self, schema: core_schema.DataclassSchema) -> JsonSchemaValue:
         cls = schema["cls"]
         fqn = f"{cls.__module__}.{cls.__name__}"
-        print("fqn: ", fqn)
+        # print("fqn: ", fqn)
         # filter attrs with leading underscore name (same behavior as in v1)
         schema["schema"]["fields"] = [f for f in schema["schema"]["fields"] if not f["name"].startswith("_")]
         json_schema = super().dataclass_schema(schema)
         properties = json_schema.get("properties", {})
         for prop, value in properties.items():
             value.pop("title", None)
-        properties.setdefault("type", {"enum": [fqn]})
+        properties.setdefault("type", {"const": fqn})
         return json_schema
 
 
-TA = TypeAdapter(FitProblem)
-schema = TA.json_schema(schema_generator=BumpsGenerateJsonSchema)
+schema = {
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$id": "refl1d-draft-02",
+}
+
+TA = TypeAdapter(FitProblem[Experiment])
+schema.update(TA.json_schema(schema_generator=BumpsGenerateJsonSchema, mode="serialization"))
 
 print(json.dumps(schema, indent=2))

refl1d/bumps_interface/fitplugin.py

@@ -8,10 +8,11 @@
 
 __all__ = ["data_view", "model_view", "new_model", "calc_errors", "show_errors"]
 
+from typing import cast
 import numpy as np
 
 from .migrations import migrate
-from ..bumps_interface.fitproblem import FitProblem
+from bumps.fitproblem import FitProblem
 from ..experiment import Experiment
 from ..sample.materialdb import air, silicon
 from ..probe.data_loaders import ncnrdata as NCNR
@@ -39,7 +40,7 @@ def load_model(filename):
     if filename.endswith(".staj") or filename.endswith(".sta"):
         from ..probe.data_loaders.stajconvert import load_mlayer
 
-        return FitProblem(load_mlayer(filename))
+        return FitProblem[Experiment](load_mlayer(filename))
         # fit_all(problem.fitness, pmp=20)
     elif filename.endswith(".zip"):
         from bumps.fitproblem import load_problem
@@ -52,11 +53,11 @@ def load_model(filename):
         with ZipFS(filename) as zf:
             for f in zf.filelist:
                 if f.filename.endswith(".py"):
-                    return load_problem(f.filename, options=options)
+                    return cast(FitProblem[Experiment], load_problem(f.filename, options=options))
     elif filename.endswith(".json"):
-        from bumps.serialize import load
+        from bumps.serialize import load_file
 
-        return load(filename)
+        return cast(FitProblem[Experiment], load_file(filename))
     else:
         return None
 

refl1d/experiment.py

@@ -27,7 +27,7 @@
     reflectivity_amplitude as reflamp,
 )
 from . import profile
-from .probe.probe import PolarizedNeutronProbe, Probe
+from .probe.probe import PolarizedNeutronProbe, Probe, QProbe, PolarizedQProbe
 from .sample import layers, material
 from .utils import asbytes
 
@@ -382,7 +382,7 @@ class Experiment(ExperimentBase):
 
     name: str
     sample: Optional[layers.Stack]
-    probe: Union[Probe, PolarizedNeutronProbe]
+    probe: Union[Probe, PolarizedNeutronProbe, QProbe, PolarizedQProbe]
     roughness_limit: float
     dz: Union[float, Literal[None]]
     dA: Union[float, Literal[None]]