feat[cartesian]: read-only data dims direct access & Fields

[FlorianDeconinck]

Description

We introduce a ROField which is a data dimensions only Field. Access is provided via .at[] which also becomes a way to access data dimensions on regular Field.

Technically, we re-use the data dims concept, empty out the cartesian axes and make sure we have protection in place against write to the ROField. Definition follows the already established data dimensions format for now: [(type, (axes...))]

Requirements

• All fixes and/or new features come with corresponding tests.
• Important design decisions have been documented in the approriate ADR inside the docs/development/ADRs/ folder.

[havogt]

Just a quick feedback (didn't look close): couldn't that be implemented with just a a normal field with no I,J or K dimension? Since our iteration space is always 3D, writing to a pure I field would also result in multiple writes from different threads to the same memory location, therefore the limit of no dimension is not really different.

[FlorianDeconinck]

Technically the ROField is exactly that: a no spatial dim, only data dim Field. I made a different FieldDescriptor to not have spaghetti code and be more verbose user side with those kind of field which are not on the grid

[FlorianDeconinck]

E.g.

class _ReadOnlyFieldDescriptorMaker(_FieldDescriptorMaker):
    def __getitem__(self, field_spec):
        if not isinstance(field_spec, collections.abc.Collection) and not len(field_spec) == 2:
            raise ValueError("ROField is defined by a tuple (type, [axes_size..])")

        dtype, data_dims = field_spec

        return _FieldDescriptor(dtype, [], data_dims)

[FlorianDeconinck]

NB: One side effect of the way I coded up .at is that you can always access data dimension that way for read operation, including on a normal Field with spatial dimensions (you get centered spatial dim)

[havogt]

I meant read-only fields and fields with no spatial dimensions are probably both useful features, but they are orthogonal. I’ll have a closer look tomorrow.

[FlorianDeconinck]

Yes, I did think about that. And technically the read-only Field is enforced by the fact you can only do .at for access which is then blocked from write at the gtscript level. So the Field itself is not RO, which is confusing for devs.

At first I wrote it as a "DataField".

Let's exchange tomorrow at community meeting

[FlorianDeconinck]

I renamed ROField to OffgridField.

I looked into using Field as-is, but Field by default have a IJK spatial axes which is used to be able to write Field[np.float64] and have it default to a 3d field.

I didn't come up with an alternate way to write the definition that was clean enough. So out of abundance of caution around the feature, I opted again for the OffgridField which forces no spatial dimensions

[havogt]

Looks good. I'll discuss with @egparedes when he is back from vacation.

[havogt]

I personally like ConstantField (constant with respect to the grid), but I think not everyone likes it for confusion with non-mutable. Here it both meanings coincide.

[FlorianDeconinck]

ConstantField would be confusing for users I think. Though, indeed, the system make it constant in stencils, the Field is still modifiable outside of gtscript

[egparedes]

I fully agree that the two features added in this PR make sense: the shortcut to access data dimensions in the current point instead of the two pairs of brackets, and a way to pass ndarrays of data which are not linked to the grid, however I disagree with the names.

Since the content of the new OffgridField concept is not linked to any position in the grid, it's not at all a physical field (https://en.wikipedia.org/wiki/Field_(physics)) and in my opinion should be just called Array, NdArray, Tensor or something along this lines. Additionally, an Array object uses absolute indexing (as any other collection instance in Python) while fields interpret indices as offsets relative to the current iteration point in the grid (numba's stencil decorator offers a similar feature by tagging a parameter as standard_indexing: https://numba.readthedocs.io/en/stable/user/stencil.html#standard-indexing). The difference with fields with data dimensions is that they are still fields, since there is a tensor quantity associated to every point in the grid.

Regarding the new .at data dimensions accessor, I would use something else since in other frameworks like jax, .at has a different meaning (https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html) and we are considering to add something similar in gt4py.next which would make the two different uses of .at in gt4py.cartesian very confusing. Other options I see could be .data , .local, .current, .this, .point, ... or single letter shortcuts in the spirit of numpy's .T like .A (array) or .D (data).

[FlorianDeconinck]

Two very good points.

The generic Array or NdArray could be confusing for users as a generic Array that transfers to any other framework (same for ndarray). I reckon we'd rather limit that concept in scope. What about GlobalMemory ? Not the greatest but hopefully you see my point.

As for the re-use of at if you are planning on shadowing jax at then of course we need something else. A single letter shortcut seems like a way to go. A for Array or Access is fine.

[egparedes]

The generic Array or NdArray could be confusing for users as a generic Array that transfers to any other framework (same for ndarray). I reckon we'd rather limit that concept in scope. What about GlobalMemory ?

Yes, I see your point and it makes sense. What about GlobalParameter or GlobalTable?

[FlorianDeconinck]

GlobalTable, I like that. It conveys both the array-ness and a form of read-only.

[egparedes]

From my side at least, you can go ahead and change the names of the two features, and I think the PR would already be ready to be merged.

[FlorianDeconinck]

Done. GlobalTable and .A for access

[FlorianDeconinck]

Ready to be merged.

[egparedes]

It looks good to me.

[egparedes]

Please update the PR to the new formatting tools (it should be very easy)

[FlorianDeconinck]

cscs-ci run

[havogt]

cscs-ci run

[egparedes]

ILGTM