Use explicit __init__ for LocalGeometry

[LiamPattinson]

Resolves (partly) #336.

Replaces general __init__ functions of the form def __init__(self, *args, **kwargs): with more explicit argument lists. Uses the same default arguments as are currently in use.

• default_inputs() functions moved to class variables DEFAULT_INPUTS, as this allows them to be used to set default values in the __init__ function and from outside the class (they're used in some GKInput classes).
• Added error checking on arrays passed to these functions to ensure they have the correct shape and size.

---

TODO:

• Update from_global_eq to be a @classmethod.
• Update from_local_geometry to be a @classmethod.
• (Guidance needed) Should normalise be called as standard by the base __init__? If so, do that. (EDIT: normalise should return a new LocalGeometry instead)

Stretch goals:

• LocalGeometry could be a @dataclass which takes R, Z, theta and b_poloidal grids as inputs and therefore represent a 'non-parameterised' local geometry. Subclasses would instead be initialised with their parameters, and would generate R, Z, theta and b_poloidal grids from there. This would resolve the problem wherein many functions defined on the base class fail due to missing data that is only generated by subclasses.
• Given the above, LocalGeometry would remove the need for the FluxSurface class, which could be removed from the project.
• Initialisation with default values could still be achieved by updating the default method to be a @classmethod. (EDIT: Not clear how this would work on the base class, probably better to avoid default behaviour where possible)
• Update docstrings throughout LocalGeometry and subclasses

[LiamPattinson]

The method from_global_eq is now a @classmethod, meaning it can be used as an alternative constructor in place of the basic __init__, and can no longer be used to modify an existing LocalGeometry.

Note

The ability to set pyro.local_geometry with a string has been removed. Rather than converting the currently loaded LocalGeometry to the desired type, this would instead replace it with an empty object, which risks allowing uninitialised data to enter the program.

This method was only used in the function load_local_geometry to pre-emptively set the type of LocalGeometry before updating it with from_global_eq. With from_global_eq being a classmethod, this is no longer needed.

This method was also used in some examples scripts, but they were easily modified to use a safer method.

(This also allowed me to delete a 2 year old FIXME comment I'd left in Pyro, which I'm quite happy about)

---

Also fixed a minor error in the Equilibrium docs while I was at it

[LiamPattinson]

The latest commit changes from_local_geometry to be a @classmethod:

# Before
local_geometry = LocalGeometryMXH() # Empty object
local_geometry.from_local_geometry(other_local_geometry)

# After
local_geometry = LocalGeometryMXH.from_local_geometry(other_local_geometry)

---

From here, I could either carry on working on this to complete the additional goals listed above, or (if you're happy with the changes) we could merge this and I'll raise a new PR once I've made some progress. I'm hoping to be done with the further goals roughly by the end of next week.

[bpatel2107]

I'm nervous about setting this to 0 as we sometimes to 1/B0...

[ZedThree]

In some ways it would make more sense to not have a default value at all for basically all of these

[LiamPattinson]

Good point, but the current default value of None is also a problem in that case (albeit one that will throw an error immediately rather than letting inf/NaN propagate throughout the code).

Is there a more sensible default value for B0? If the current defaults are pulled from a standard test case, is B0 defined there?

[LiamPattinson]

@ZedThree I'd like to eventually get these classes to a point where none of the inputs are optional (besides things like delta in LocalGeomtryMXH), but the option to build one with sensible defaults could still be achieved with a class method, e.g. LocalGeometry.default(), as that would make testing a lot easier.

[bpatel2107]

I think we usually have checks for B0 usually like

if local_geometry.B0:
    beta = 1 / local_geometry.B0**2
else:
    beta = 0.0

Hence the default of None as having beta=0.0 by default is safest.

[bpatel2107]

This looks great, makes LocalGeometry more like the other classes and I think being consistent overall is a good idea. Making it a @dataclass seems likea good way to go ahead.

normalise adds/converts the units of the different LocalGeometry parameters. When all is said and done they should have units. The issue becomes whether or not the units or way to convert units is known when initialising the object. When loading from Equilibrium it is straightforward, but when loading from input files there may be times where the units are not immediately obvious without looking at other parts of the input file. This could be handled in the specific code such that from_gk_data includes the appropriate units/convention (should be the called load_from_gk_data?). So normalise should be in the __init__ but we need to make sure we know the appropriate units.

[LiamPattinson]

Finally managed to get somewhere with this refactor, though it looks like I took so long that I now have a really messy merge on my hands...

LocalGeometry is now a dataclass, and local geometries are now always built using either their own __init__, from_global_eq, or from_local_geometry. They have a consistent set of attributes whichever method was used, and no further set up is required after initialisation (e.g. no need to assign bunit_over_b0 after building the object).

The fitting process for shaping coefficients has been modified to use @classmethod's and @staticmethod's throughout. This means that self is never referenced during the fitting process, and all necessary data is determined before initialisation (see LocalGeometry._init_with_shape_params). I had to add some extra machinery to get this to work in a generic manner, such as the ShapeParams classes and _shape_params property. This reduces the amount of repeat code in subclasses, at the cost of a more complex base class.

There are still some tasks to complete, two of which I've listed in comments. Further work includes:

• Clean up docs in subclasses
• Remove defaults from subclass __init__ functions. This may be challenging, as quite a lot of the code and the tests currently depend on default behaviour.
• from_global_eq should not modify the normalisation -- the responsibility for this step should be handed to Pyro.
• normalise should return a new local geometry

[LiamPattinson]

When initialising a local geometry, the code was previously overwriting dpsidr if it was provided via a GK input file, but not if it was determined via an equilibrium or another local geometry. This feels like a bug.

This check is included as otherwise a few tests fail that I can't verify for myself. It also necessitates an overwrite_dpsidr variable in all of the subclass __init__ functions. If this is unintended behaviour, I can remove it throughout.

[bpatel2107]

dpsidr should be determined from LocalGeometry parameters alone, you never actually specify it as an input to a code. You can of course get it from a full equilibrium so that could be a good check but in theory you should always take the one from the LocalGeometry parameters, its like a fit I guess. I guess it should always be overwritten.

[LiamPattinson]

I've been playing around with this and I ran into a bit of a units issue. If we build from a GK input file, then everything is just a float until we attach units during the normalisation step, after which dpsidr is assigned units of lref * bref. However, if we calculate it from an equilibrium file, dpsidr will have units of just [length], as rho * bunit_over_b0 / q has units [length] [dimensionless] / [dimensionless], and when we try to normalise this to lref * bref it throws an error. I think we need to use the provided dpsidr when building from an equilibrium, but otherwise should be fine to recalculate.

[LiamPattinson]

I did some further playing around and it seems this issue appears when building from another local geometry too. It looks like the best solution is to use dpsidr if it's available, and only calculate it as a backup. I've come up with a cleaner solution to the problem though, so it's not a big deal -- not sure why I thought I needed overwrite_dpsidr, as it seems just using a default value of None is sufficient.

[bpatel2107]

The issue here with bunit_over_b0 is that it technically is Bunit with units of B0 so when bref = B0 then it dimensionally makes sense. To be honest it may be worth adding in the correct units for Bunit which is a just a magnetic field rather than having Bunit/B0.

[LiamPattinson]

I had to disable this test and a few others for FourierGENE/CGYRO, as I think they were comparing values calculated on different theta grids, and I can't figure out how the test parameters were calculated. Where the tests say the data was determined using sympy, it would be good if the code used to generate the data could be included in the docstrings.

[bpatel2107]

These are analytic solutions for grad_r in a simplified geometry so they should always be true. We can look at getting the sympy version written up somewhere

[bpatel2107]

Example code using sympy to get the golden answers. Given that we fit say a FourierGENE to a Miller fit we expect it to be close but not exact. Maybe the tolerance should be increased though nothing here should change an answer that already worked... Sometimes we need to be careful with integrals used in minimising b_poloidal and how they are set up with boundaries in theta Maybe the theta grid has changed which could tweak the integral and modify the fit

from sympy import Symbol, simplify, cos, sin, asin, sqrt


def simplify_grad_r(rho, Rmaj, shift, kappa, s_kappa, delta, s_delta):

    theta = Symbol("theta")

    # Not used but nice to see
    R = Rmaj + rho * (cos(theta + asin(delta) * sin(theta)))
    Z = rho * kappa * sin(theta)

    dRdr = (
        shift
        + cos(theta + asin(delta) * sin(theta))
        - sin(theta + asin(delta) * sin(theta)) * sin(theta) * s_delta
    )

    dRdtheta = (
        -rho * sin(theta + asin(delta) * sin(theta)) * (1 + asin(delta) * cos(theta))
    )

    dZdr = kappa * sin(theta) + s_kappa * kappa * sin(theta)
    dZdtheta = kappa * rho * cos(theta)
    g_tt = dRdtheta**2 + dZdtheta**2

    grad_r = sqrt(g_tt) / (dRdr * dZdtheta - dRdtheta * dZdr)

    return simplify(grad_r)


# Simple circular plasma
rho = 0.5

Rmaj = 3.0
shift = 0.0

kappa = 1.0
s_kappa = 0.0

delta = 0.0
s_delta = 0.0

my_grad_r = simplify_grad_r(rho, Rmaj, shift, kappa, s_kappa, delta, s_delta)
print(f"Grad_r for circular plasma: {my_grad_r}")


# Elongation gradient
rho = 0.5

Rmaj = 3.0
shift = 0.0

kappa = 1.0
s_kappa = 1.0

delta = 0.0
s_delta = 0.0

my_grad_r = simplify_grad_r(rho, Rmaj, shift, kappa, s_kappa, delta, s_delta)
print(f"Grad_r for circular plasma with s_kappa: {my_grad_r}")


# High shaping
rho = 0.5

Rmaj = 3.0
shift = 0.1

kappa = 2.0
s_kappa = 1.0

delta = 0.5
s_delta = 0.2

my_grad_r = simplify_grad_r(rho, Rmaj, shift, kappa, s_kappa, delta, s_delta)
print(f"Grad_r for shaped plasma: {my_grad_r}")

Running this code results in

Grad_r for circular plasma: 1.00000000000000
Grad_r for circular plasma with s_kappa: 1.0/(sin(theta)**2 + 1)
Grad_r for shaped plasma: sqrt(0.25*(0.523598775598299*cos(theta) + 1.0)**2*sin(theta + 0.523598775598299*sin(theta))**2 + 1.0*cos(theta)**2)/(2.0*(0.523598775598299*cos(theta) + 1.0)*sin(theta)*sin(theta + 0.523598775598299*sin(theta)) + 1.0*(-0.2*sin(theta)*sin(theta + 0.523598775598299*sin(theta)) + cos(theta + 0.523598775598299*sin(theta)) + 0.1)*cos(theta))

[LiamPattinson]

Thanks for this, I tried adapting the code to also generate R and Z to make sure we were comparing the same flux surfaces. It looks like we are, but each value of theta corresponds to a different point along the flux surface:

It's a bit too late on a Friday for me to figure out what's going on here. Is it the case that the definition of theta is different in the Miller geometry compared to the Fourier geometries?

[bpatel2107]

Yes the definition is different which makes a lot of this quite annoying. You need to determine theta when given an array of R and Z. You can see the definitions would be different when looking at the parameterisation for each type.

[LiamPattinson]

I think I have this resolved in the latest commit, but it's a bit hacky! I tried to find an analytic way to convert between the two theta systems, but in the end gave up and just added a small routine that takes R, Z, and theta from one local geometry and finds which Miller-theta should correspond to that using a least squares fit.

[LiamPattinson]

I've finished merging and attempted to recreate the new GENE features using my new methods, but I can't tell if I've done so correctly or not. This section of the code looks like it's untested, and I don't really know what it's doing.

[bpatel2107]

Essentially the full data needed to populate LocalGeometry is contained inside a different file which requires a bit of post-processing to get the data in a form that can actually be read by LocalGeometry. What you do here looks to be correct. This is tested here

pyrokinetics/tests/local_geometry/test_fourier.py

Line 383 in 765d0f1

def test_tracer_efit_eqdsk():

[LiamPattinson]

That test doesn't seem to touch the new code in gk_code/gene.py, so I'm not fully confident that my changes haven't broken something. I can leave it alone for now since it's slightly outside the scope of the PR, but it would be good to get a test in that uses one of those new geometry files alongside a GENE input file directly.

[bpatel2107]

Right I see what you mean! I just fit a FourierGENE to the GEQDSK and compare to a new geometry file I made independently. I can add that in but yes its outside the scope of this PR

[LiamPattinson]

With the aim of promoting immutability, normalise() now returns a new object, and from_global_eq no longer updates the current normalisation as a side effect. In order to merge LocalGeometry and FluxSurface, it would be ideal to remove the norms argument from from_global_eq entirely, so I'll have a think about how best to achieve that. I also still have to fix a couple of broken tests and go over the docs.

[LiamPattinson]

I think that's about as far as I'd like to take this PR. I'd still like to merge FluxSurface and LocalGeometry, but it's less essential. I've reworked the docs, removed most default __init__ parameters, and fixed the broken tests. I also removed the norms argument from from_global_eq, as if I do get around to merging FluxSurface and LocalGeometry, it would be good to be able to do so without creating a Normalisation first.