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.

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...

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.

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)

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).

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

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