Magnetic Diagnostics API

[dpanici]

Adds new class AbstractDiagnostic which is a new base class for synthetic diagnostic objects.

• AbstractDiagnostic realizations must

  • define the points at which they may require the total magnetic field (as determined by Bplasma + Bcoil if outside of the plasma
    • TODO: Might need a separate class for objects which need the INTERNAL magnetic field, but I actually think the current API would work as the internal magnetic field can be computed as like auxiliary data with the _compute_data method already in this
  • as well as implement a compute method for computing the signal they represent
  • implement a _compute method which takes in the magnetic field (or vector potential) at the points they require (for e.g. integration over a rogowski coil), as well as a dict of aux data
  • a _compute_data method for compiting this aux data from eq_params, field_params as well as its owndiag_params (here for example a RogowskiCoil computes its geometric information. A ThomsonScattering may here compute the intersection of the chord and the plasma, as another example)
  • may have DOFs to be optimized
    Current realizations: RogowskiCoilXXX for each coil type, PointBMeasurements for ptwise B measurements like magnetic probes.

• DiagnosticSet objects are collections of AbstractDiagnostic realizations

  • mainly operating like MixedCoilSet for now, its job is to collect the diagnostics in one object,
  • collect and concatenate the points at which B is needed for them so that in MeasurementError it can be computed efficiently

TODO

• How to easiest pass in target/bounds?
• Maybe implement a DiagnosticSet that is like a CoilSet too?
• have diagnostics have two eval_x attrs required, one for eval_B and other for eval_A and update MeasurementError to compute both A and B at these points
• maybe move the computation of magnetic field to the DiagnosticSet, since some diagnostics may not need the coil field, for instance, if they are internal measurements
• currently need _diagnostics to be a static attr to avoid JAX issues with MeasurementError but not sure why, maybe because we are using its method? but that should be fine I thought... so I dont think that is the issue
• Add method of specifying correlated uncertainties (i.e. generalize W weight matrix from diagonally as it implicitly is now to arbitrary W) within a single diagnostic (like here within a single diagnostic objective)
  • for MeasurementError, allow weights to be a matrix and assume it to be the inverse of the covariance, and Cholesky factor it and make _scale method for this class using that factorization
• merge this into Add point B measurement error objective #1436 so tests can run
• Add flux loops from branch which has those implemented
  • make them and Rogowski from same parent class

Resolves #1783

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    5.13 %    |     3.836e+03      |     4.033e+03      |    196.80    |       36.55        |       31.33        |
  test_proximal_jac_w7x_with_eq_update   |   -0.57 %    |     6.799e+03      |     6.760e+03      |    -38.56    |       161.42       |       159.80       |
  test_proximal_freeb_jac                |   -0.18 %    |     1.317e+04      |     1.315e+04      |    -24.30    |       78.62        |       78.32        |
  test_proximal_freeb_jac_blocked        |   -0.61 %    |     7.640e+03      |     7.593e+03      |    -46.61    |       67.77        |       69.78        |
  test_proximal_freeb_jac_batched        |   -0.06 %    |     7.599e+03      |     7.595e+03      |    -4.29     |       68.28        |       69.39        |
  test_proximal_jac_ripple               |   -1.73 %    |     7.619e+03      |     7.487e+03      |   -131.84    |       69.21        |       70.22        |
  test_proximal_jac_ripple_spline        |   -1.05 %    |     3.481e+03      |     3.444e+03      |    -36.45    |       70.70        |       73.12        |
  test_eq_solve                          |   -2.02 %    |     2.062e+03      |     2.021e+03      |    -41.73    |       124.94       |       125.97       |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[f0uriest]

Is there still a plan to have classes for different diagnostics? I didn't realize those would each be objectives. I thought it would be something like

class AbstractDiagonstic(ABC, Optimizeable?):
    def compute_measurement(self, eq, field):
        ...
    
    def _compute_from_B(self, B):
        ...
class RogowskiCoil(AbstractDiagnostic):
    ...

class MeasurementError(_Objective):

    def __init__(self, eq, field, diagnostic, target, weight):
        ...

    def compute(self, params):
        B = self.field.compute_magnetic_field(...)
        return self.diagnostic._compute_from_B(B)

I like the idea of separate classes for different diagnostics, independent from the Objective API for doing analysis etc. The diagnostic only knows about how to compute a synthetic measurement, but nothing about the expected value from experiment. The objective then couples the diagonstic class with the target/bounds/uncertainty from experiment to compute a residual/z score.

Another thing to consider is uncertainties. We can handle simple uncorrelated uncertainty with weight=1/uncertainty but often there are off diagonal terms in the covariance matrix that are important. This might require the Measurement objective to have special logic in compute_scaled_error etc.

At some point we may also want to consider optimizing the placement of diagnostics, this is something that a lot of experimentalists think about and could be a good use of DESC. So then AbstractDiagnostic would also inherit from Optimizeable

[f0uriest]

_all_eval_x_rpz might be better as an abstract property/attribute itself?

Also, the variable name implies it but should document somewhere that x here is in R,phi,Z. And that the compute function should expect B in R,phi,Z basis as well

[dpanici]

I can do that, but it seems that that way does not explicitly require that the attributes themselves be set, only that the properties exist?

Like one still could instantiate a subclass with the property _all_eval_x_rpz but never fill it?

[f0uriest]

can't this just be a class method?

[dpanici]

I can't import it without getting a circular import so I have to import it locally. I could do that as a class method I guess

[dpanici]

Is there still a plan to have classes for different diagnostics? I didn't realize those would each be objectives. I thought it would be something like

class AbstractDiagonstic(ABC, Optimizeable?):
    def compute_measurement(self, eq, field):
        ...
    
    def _compute_from_B(self, B):
        ...
class RogowskiCoil(AbstractDiagnostic):
    ...

class MeasurementError(_Objective):

    def __init__(self, eq, field, diagnostic, target, weight):
        ...

    def compute(self, params):
        B = self.field.compute_magnetic_field(...)
        return self.diagnostic._compute_from_B(B)

I like the idea of separate classes for different diagnostics, independent from the Objective API for doing analysis etc. The diagnostic only knows about how to compute a synthetic measurement, but nothing about the expected value from experiment. The objective then couples the diagonstic class with the target/bounds/uncertainty from experiment to compute a residual/z score.

Another thing to consider is uncertainties. We can handle simple uncorrelated uncertainty with weight=1/uncertainty but often there are off diagonal terms in the covariance matrix that are important. This might require the Measurement objective to have special logic in compute_scaled_error etc.

At some point we may also want to consider optimizing the placement of diagnostics, this is something that a lot of experimentalists think about and could be a good use of DESC. So then AbstractDiagnostic would also inherit from Optimizeable

These are all good points, I mainly shied away from the optmizable diagnostics part because it is gonna add a lot more hassle on my part and I just wanted to get something working for my project...

[f0uriest]

Yeah, I'm definitely not saying we have to worry about optimizing sensor placement or dealing with full covariance matrices now, but I just want to make sure we give it a bit of thought so that we don't design ourselves into a corner that will require major API changes to add those features later.

I think broadly we could copy the API from MagneticField and Coil but sort of replace compute_magnetic_field(x) with compute_measurement(eq, field, ...)
We could also have them inherit from Optimizable but for now don't give them any optimizable DoFs, or add a diagnostic_fixed flag to the objectives that we hard code to True for now (similar to what we did for the boundary error objective before single stage)

For correlated uncertainty, if its only correlated within a single objective that could be handled by just allowing weight to be a matrix (ie cholesky factor of the inverse covariance, though again we don't need to deal with supporting this now since its effectively a backwards compatible change). If we want to allow for coupled uncertainty between different objectives that might be trickier so might be worth thinking it through.

[dpanici]

API Issues I am wrestling with currently:

• if diagnostic positions are optimizable (or potentially optimizable), then the whole vectrization method I came up with in this PR won't work as is (as in this PR I in the build of MagneticDiagnostics assemble all the points I need to eval B at). I guess this could instead happen in the compute method of said super objective, and it would have to
  • take in the diag_params of each sub-diagnostic, which I also don't know how we will do right now since it is an arbitrary number, maybe like how I did it in Add ShareParameters Linear Objective #1320, but ti would also need to know about eq_params and field_params, then an arbitrary number of diag_params...
  • then also somehow know from each diag_params what optimizable DOF it needs to use to get the x positions it needs to evaluate B at (this is complex, as in this PR it would simply just be diag_params["R"] etc for the pointB measurement, but in future we want flux loops, which are themselves arbitrary coil types. I guess this could be alleviated if we also add diagnostics to the compute index (have loops diags subclass from coils, and add any new ones which dont have obvious analogs to subclass from like this pointBMeasurement one)

[f0uriest]

One possible idea would be to add something like a CoilSet but for diagnostics, like DiagnosticSet or something. And then the reconstruction objective requires a single DiagnosticSet rather than individual diagonstics? Then all the vectorization can be done there (similar to how CoilSet uses a more efficient way of computing the total field).

[dpanici]

Is there still a plan to have classes for different diagnostics? I didn't realize those would each be objectives. I thought it would be something like

class AbstractDiagonstic(ABC, Optimizeable?):
    def compute_measurement(self, eq, field):
        ...
    
    def _compute_from_B(self, B):
        ...
class RogowskiCoil(AbstractDiagnostic):
    ...

class MeasurementError(_Objective):

    def __init__(self, eq, field, diagnostic, target, weight):
        ...

    def compute(self, params):
        B = self.field.compute_magnetic_field(...)
        return self.diagnostic._compute_from_B(B)

I like the idea of separate classes for different diagnostics, independent from the Objective API for doing analysis etc. The diagnostic only knows about how to compute a synthetic measurement, but nothing about the expected value from experiment. The objective then couples the diagonstic class with the target/bounds/uncertainty from experiment to compute a residual/z score.
Another thing to consider is uncertainties. We can handle simple uncorrelated uncertainty with weight=1/uncertainty but often there are off diagonal terms in the covariance matrix that are important. This might require the Measurement objective to have special logic in compute_scaled_error etc.
At some point we may also want to consider optimizing the placement of diagnostics, this is something that a lot of experimentalists think about and could be a good use of DESC. So then AbstractDiagnostic would also inherit from Optimizeable

These are all good points, I mainly shied away from the optmizable diagnostics part because it is gonna add a lot more hassle on my part and I just wanted to get something working for my project...

Doing something like this might involve making many redundant classes just to have for example a RogowskiCoil for each curve type: RogowskiCoilFourierXYZ, RogowskiCoilSplineXYZ etc. I can not think of a way to have separate diagnostics which are also optimizable and avoid this

[dpanici]

For correlated measurements, we should be able to use most of our existing machinery if we just make a new objective called something like MeasurementError which takes in all of the diagnostics, and the passed-in weights can be either a single numpy array (corresponding to a diagonal covariance matrix, where the weights passed in are the inverse of the square root of the diagonal variances), or it can be a (positive def, symmetric) matrix W which is assumed to be the inverse of the (also positive def, symmetric) covariance matrix $\Omega$ i.e. $W = \Omega^-1$.

Since these are pos def symmetric, we can do Cholesky to factor
$$W = AA^T $$
and store that when we build the objective.

Then we can modify its compute_scaled (or maybe its _scale or whatever) to return
$A^T f$
so that the overall sum of squares in the compute_scalar becomes the generalized least squares problem
$\mathbf{f}^T W \mathbf{f} = \mathbf{f}^T AA^T \mathbf{f} $

where $\mathbf{f}$ is the measurement residual returned by the MeasurementError compute. This also should allow us to still use Gauss-Newton Hessian approximations like we do in lsq-exact