Skip to content
Jim Pivarski edited this page Mar 3, 2021 · 8 revisions

Design plans, as they evolve

This is a follow-up on Discussion #30 with a general outline of what I'm implementing and why.

The basic objects are:

  • Vectors: 2D, 3D, and 4D "directions (with magnitude)" in space. This is essentially a 2, 3, or 4 component object.
  • Points: 2D, 3D, and 4D "positions" in space. This is also a 2, 3, or 4 component object, differentiated only in how Transforms and FrameTransforms act on it.
  • Transforms: 2D → 2D, 3D → 3D, and 4D → 4D changes in direction and magnitude. The General2DTransform, General3DTransform, and General4DTransform have 2², 3², or 4² components, but PlanarTransform, SpatialTransform, and LorentzTransform (see below) have only 1, 3, and 6 components, generating the rest of the matrix algorithmically. The mathematical functions only assume that xx, xy, yx, yy, etc. exist and call them, be they attributes or properties.
  • FrameTransforms: 2D → 2D, 3D → 3D, and 4D → 4D changes in position, direction, and magnitude. This is a Transform and a Vector for rotation/boost and displacement.

There are two forms for each of the above:

  • a single object, representing a single Vector, Point, Transform, or FrameTransform
  • an array of such objects, named VectorArray, PointArray, TransformArray, or FrameTransformArray.

The names for each number of dimensions are:

  • 2D: Planar, as in PlanarVector, PlanarPoint, PlanarTransform (1 component), PlanarFrameTransform (3 components), General2DTransform (4 components), General2DFrameTransform (6 components)
  • 3D: Spatial (like the above)
  • 4D: Lorentz (I know that the Lorentz-ness is defined by how it's transformed, not in the vector itself, but we'll use this word because Minkowski space is more important than 4D Euclidean space, just as 2D and 3D Euclidean spaces are more important than any other metrics).

Vectors and Points are distinguished as classes so that they can have different __add__ and __sub__ methods (etc.). Initially, only Vectors will actually be implemented (and hence no FrameTransforms, either). We just want to establish that the name exists, so that we know what to call this generalization later, and how it fits into a class hierarchy (i.e. we'll create a mixin class named "Vector", rather than just assuming that everything is a vector, so that "Point" can be added without changing the existing interface for vector-users).

The storage class representing both vectors and points consists of three objects, which have independent coordinate systems:

name applies to coordinate systems
azimuthal 2D, 3D, 4D xy, rphi (two members each; x, y Cartesian and r, phi polar)
longitudinal 3D, 4D z, theta, eta, w (one member each; z is Cartesian, theta is dip angle, eta is like pseudorapidity and w is like rapidity)
temporal 4D t, tau (one member each; t is like time/energy and tau is like proper time/mass)

Computing coordinates in another system sometimes requires other pieces of the vector. For instance, to get theta from a z longitudinal, we need the azimuthal (specifically, we need sqrt(x**2 + y**2) if xy or r if rphi). The math functions in compute are specialized for combinations of coordinate systems, including those math functions that compute components of a vector: e.g. a function to compute phi has a version that works with azimuthal xy (by computing the atan2) and a version that works with azimuthal rphi (by passing through the phi).

That combinatorically increases the number of functions, but that's what we were doing anyway in the old libraries. At least we should acknowledge it and name the functions appropriately. Splitting azimuthal, longitudinal, and temporal into pieces also increases the number of combinations, but not all combinations need to be specialized: some can internally switch to a more convenient set of coordinates and call another. (This will also make it easier to find the hotspots for future optimizations because they'll be explicitly converting and calling other functions, rather than doing things in place.) Also, not all three pieces—azimuthal, longitudinal, and temporal—will be needed for every function. If we defined coordinate systems in a non-piecemeal way, then somebody wanting "rphi, w, and tau" would have to implement a whole suite of functions if it wasn't already there.

The compute module

All of the mathematical calculations take place in flat functions in a module named compute. The compute functions take numerical or array (not class) arguments, have no branches (if statements or loops) or assignments in the bodies of the functions, and maybe should consist entirely of ufuncs. They are to be duck-typed into a variety of contexts, so that the mathematics is not duplicated. They can be tested against the corresponding ROOT Math::Lorentz classes once and for all, so that we're not uncertain about the mathematics.

These functions will be organized as:

class whatitproduces(Binary):   # or Unary, Trinary, ...
    dispatch_map = {
        (A1, L1, T1, A2, L1, T1): whatitproduces.A1L1T1_A2L2T2,
        ...
    }

    @staticmethod
    def dispatch(lib, v1, v2, ...):
        return whatitproduces.dispatch_map[
            type(v1.azimuthal), type(v1.longitudinal), type(v1.temporal),
            type(v1.azimuthal), type(v1.longitudinal), type(v1.temporal),
        ](lib, *(
            v1.azimuthal.coordinates + v1.longitudinal.coordinates + v1.temporal.coordinates + 
            v2.azimuthal.coordinates + v2.longitudinal.coordinates + v2.temporal.coordinates
        ))

    @staticmethod
    def A1L1T1_A2L2T2(lib, x1, y1, z1, t1, x2, y2, z2, t2, ...):
        # compute

The whatitproduces can be multiple quantities, separated by underscores, in which case, all functions in the class return a tuple of results. The lib is a numpy module equivalent, such as numpy, cupy, jax.numpy, etc. This is what gets used to fetch ufuncs. The remaining arguments are the one, two, or (unlikely) more arguments that may be individual vectors or arrays of vectors, with no opinion about their backend. If a calculation doesn't require all three types of coordinates, only the required ones go in the dispatch_map and are tested in dispatch. (That will be uniform across a class.) What I've denoted "A1L1T1_A2L2T2" here are the combinations of coordinate systems. These need to be exhaustive for the types of coordinates used (as many as 16). Any missing definitions would be caught by Unary, Binary, Trinary, ... protocols. Non-vector arguments don't count toward this number, but the non-vector could potentially be an array. Another kind of non-vector is a Transform, but these aren't dispatched by coordinate system either: they have xx1, xy1, yx1, yy1, etc. arguments.

All of the functions in these classes are static methods—the class is only used as a bookkeeping mechanism (to simplify naming conventions). Thus, these functions simply take a library and arguments that may be scalar numbers or arrays of numbers (or jagged arrays without record structures). That makes them very portable to all the backends, NumExpr (if we capture its string through inspect and drop the "lib."), Numba, etc.

Or maybe instead of "classes that aren't classes" (just naming conventions), maybe each should be a submodule in compute. Then there can be a copy-paste template for getting started...

Clone this wiki locally