pmelchior / scarlet2

Scarlet, all new and shiny
MIT License
15 stars 4 forks source link

New Parameter specification #56

Closed pmelchior closed 5 months ago

pmelchior commented 5 months ago

this PR follows from #55

The new interface splits creating and initializing the model from specifying optimization/sampling parameter. An example is below:


with Scene(model_frame) as scene:

    for center in centers:
        # all inputs to models are int, float, complex or jnp arrays
        center = jnp.array(center)
        spectrum = 100 * jnp.ones(5)
        morph = jnp.ones((21, 21))

        Source(
            center,
            ArraySpectrum(spectrum),
            ArrayMorphology(morph)
        )

# create a lookup structure how to find a specific node of the equinox model scene so that it can be optimized
# can find variable names automatically, but better to give names explicitly
parameters = scene.make_parameters()
for i in range(len(scene.sources)):
    parameters += Parameter(scene.sources[i].spectrum.data, name=f"spectrum.{i}", constraint=constraints.positive, stepsize=spec_step)
    parameters += Parameter(scene.sources[i].morphology.data, name=f"morph.{i}", constraint=constraints.positive, stepsize=0.1)

# can remove parameters by name
parameters -= "spectrum.4"