firedrakeproject / asQ

A library for implementing ParaDIAG timestepping algorithms using Firedrake
MIT License
4 stars 1 forks source link

Allatonce subfunctions (waiting on #136) #141

Closed JHopeCollins closed 11 months ago

JHopeCollins commented 11 months ago

This PR improves how we access the solution at individual timesteps of an AllAtOnceFunction.

Currently these are accessed via the get_field and set_field methods (and the analogous get_component and set_component methods for accessing one component of a timestep when using mixed function spaces). These are a bit clunky and not really a very pythonic style.

Instead, we can create a set of Functions to view the solution at each timestep. This is analogous to how Function.subfunctions works for accessing individual components of a mixed Function. Each subfunction is a Function in its own right that is a view over the data in the full mixed Function.

For the AllAtOnceFunction, we follow a very similar pattern. We create a tuple of Functions, (AllAtOnceFunction._fields) each of which views the data for a particular timestep in the MixedFunctionSpace that represents the local slice of the timeseries. When the FunctionSpace for a single timestep is mixed, the Functions in AllAtOnceFunction._fields are also mixed.

The Functions that view each timestep are accessed by indexing the AllAtOnceFunction. For example, previously we would have written:

u = fd.Function(aaofunc.field_function_space)
aaofunc.get_field(i, uout=u)

Now we can just write:

u = fd.Function(aaofunc.field_function_space)
u.assign(aaofunc[i])

which is much closer to standard Firdrake syntax. To access a particular component (assuming aaofunc.field_function_space is mixed):

aaofunc[i].subfunctions[0].assign(u.subfunctions[0])

(I'm not completely attached to accessing these Functions by directly indexing the AllAtOnceFunction. We could also name them directly like for firedrake.Function.subfunctions. e.g. we could have AllAtOnceFunction.fields and access them via aaofunc.fields[i]. Happy to discuss.)

Another change in this PR is that AllAtOnceFunction.transform_index will no longer give you the index for a particular component in the mixed space for the local slice of the timeseries. Getting this index is kind of an implementation detail (specific to the fact that we've implemented each slice as a mixed space) so shouldn't be necessary to use from user-land. There is now the "internal" method AllAtOnceFunction._component_indices for when we need to access these indices inside asQ.

The changes to the examples and case_studies scripts show how these changes are used. I think it makes things simpler/nicer to read but please let me know what you think!

JHopeCollins commented 11 months ago

This is waiting on #136 because the uses of get_field in the vertical slice scripts need removing, but that PR changes the vertical slice scripts substantially. Once that PR is merged into master I'll merge it in here to update those scripts.

JHopeCollins commented 11 months ago

Yes, I prefer it this way!

We will need to warn Hiroe that this has changed a bit.

Great, thanks! I've just posted in the asQ slack about it.