Improvement of response matrices

[lfarv]

This PR implements the features requested in #962.

• Fix the use of custom variables,
• allow to instantiate a ResponseMatrix with ring=None (the default), if no lattice is necessary for any variable or observable,
• improve the error message if trying to access a RefptsVariable with ring=None,
• ensure that all variables are restored if an exception occurs during building the matrix,
• adds a hook called after each variable variation

[JeanLucPons]

Hello,

I'm still having this issue:

  File "/operation/common/miniconda/envs/jlp-py312/lib/python3.12/site-packages/at/latticetools/response_matrix.py", line 445, in __init__
    var.refpts = ring.get_uint32_index(var.refpts)

AttributeError: 'NoneType' object has no attribute 'get_uint32_index'

I managed to make it work by doing the mod below:

diff --git a/pyat/at/latticetools/response_matrix.py b/pyat/at/latticetools/response_matrix.py
index d0dd2287..f3e09205 100644
--- a/pyat/at/latticetools/response_matrix.py
+++ b/pyat/at/latticetools/response_matrix.py
@@ -441,8 +441,10 @@ class ResponseMatrix(_SvdSolver):
                 beg = end
 
         # for efficiency of parallel computation, the variable's refpts must be integer
-        for var in variables:
-            var.refpts = ring.get_uint32_index(var.refpts)
+        if ring is not None:
+            for var in variables:
+                if hasattr(var,"refpts"):
+                    var.refpts = ring.get_uint32_index(var.refpts)
         self.ring = ring
         self.variables = variables
         self.observables = observables

I expect that the response matrix engine can be fully ring independent. I didn't find a way to define a CustomObservable without a ring ? I would expect a CustomObservable that can accept free input objects as for CustomVariable

Here is a sample code and I would like that the ring is not passed at all to AT functions:

def set_int_field(value:float,e:Magnet,ring:at.Lattice|None):
    e.set_strength(value)

def get_int_field(e:Magnet,ring:at.Lattice|None):
    return e.get_strength()

def eval_tune(ring):  # Here ring should be a free param as Magnet object for a CustomVariable
    return ring.get_tune()

def get_tune_matrix(ring:at.Lattice,qdTune:list[Magnet],qfTune:list[Magnet]):

    variables = None
    DT = 1e-5

    varlist = []
    for q in qdTune:
        v = CustomVariable(set_int_field, get_int_field, q, delta=DT)
        varlist.append(v)
    for q in qfTune:
        v = CustomVariable(set_int_field, get_int_field, q, delta=DT)
        varlist.append(v)
    variables = VariableList(varlist)

    observables = ObservableList([RingObservable(eval_tune)])
    observables.evaluate(ring=ring)
    tune_response = ResponseMatrix(variables, observables, ring)
    tune_response.build(use_mp=True)
    tune_response.solve()

    mat = tune_response.correction_matrix()
    return mat

[lfarv]

@JeanLucPons:

I did some other corrections to avoid crashes with ring = None.

I didn't find a way to define a CustomObservable without a ring ?

For "custom observables", the easiest is to use the base Observable class.

With this, here is an updated version of your example:

def set_int_field(value: float, e:Magnet, **_):
    e.set_strength(value)

def get_int_field(e: Magnet, **_):
    return e.get_strength()

def eval_tune(eval_data, **_):
    return eval_data.get_tune()

def get_tune_matrix(eval_data, qdTune: list[Magnet], qfTune: list[Magnet]):

    DT = 1e-5

    variables = at.VariableList()
    for q in qdTune:
        variables.append(CustomVariable(set_int_field, get_int_field, q, delta=DT))
    for q in qfTune:
        variables.append(CustomVariable(set_int_field, get_int_field, q, delta=DT))

    obs = Observable(eval_tune, eval_data, name="tunes")
    observables = ObservableList([obs])
    observables.evaluate()
    tune_response = ResponseMatrix(variables, observables)
    tune_response.build(use_mp=True)
    tune_response.solve()

    mat = tune_response.correction_matrix()
    return mat

Is this what you expect ? I am still adding a few features, let me know if you have ideas !

[lfarv]

I added the possibility to add keywords to the VariableBase.set and VariableBase.get methods. They are passed to the low-level getfun and setfun.

This gives the possibility to give "modifiers" to these functions (verbose, for instance).

These keywords given to set() or get() add or replace the ones given in the variable constructor. So it is advisable that the custom getfun and setfun functions accept, and usually ignore these keywords, unless they want to make use of them:

def set_int_field(value: float, e:Magnet, **_):
    # Keywords are ignored
    e.set_strength(value)

def get_int_field(e: Magnet, **_):
    # Keywords are ignored
    return e.get_strength()

v = CustomVariable(set_int_field, get_int_field, q, delta=DT)
v.get(verbose=True)  # get_int_field receives and ignores the "verbose" keyword

Similarly, there are optional keywords in ObservableList.evaluate.

def eval_tune(eval_data, scaling=1.0, **_):
    # Accept a scaling custom keyword
    return scaling*eval_data.get_tune()

obs = Observable(eval_tune, eval_data, name="tunes")
observables = ObservableList([obs])
observables.evaluate(scaling=0.9). # eval_tune gets the "scaling" value

[JeanLucPons]

Thanks @lfarv
It works as I except for a generic response matrix.

[lfarv]

@swhite2401:

I find it very annoying to have to input ring every time an evaluate/get/set function is called,

The modifications described above give a solution for that: when adding a ring keyword when instantiating a RefptsVariable, it will be stored and used as a "default ring" when using .get() or set(). However, a ring specified in get() or set() will have priority, ensuring that the previous behaviour is preserved. The same is true for observables. Example:

v = RefptsVariable("QD2[AE]", "PolynomB", index=1, ring=ring)
v.get()
np.float64(-2.6720252472931687)

Here ring is a special case: any keyword at instantiation will be passed to the evaluation function.

[lfarv]

The problem with the tests comes from the release of a new version of h5py. PyAT works correctly, we just have to wait for a correction. The failure happens for old macOS versions, everything else is fine.

[lfarv]

@swhite2401 and @JeanLucPons : as far as I can see, this version answers all the points discussed in #962 and here above. Is there any remaining point, or can we merge this one, and take care of the next points in another PR?