Warning

You are reading a version of the website built against the unstable main branch. This content is liable to change without notice and may be inappropriate for your use case. You can find the documentation for the current stable release here.

firedrake.ensemble package

Submodules

firedrake.ensemble.ensemble module

class firedrake.ensemble.ensemble.Ensemble(comm: Comm, M: int, **kwargs)[source]

Bases: object

Create a set of space and ensemble subcommunicators.

Wrapper methods around many MPI communication functions are provided for sending Function and Cofunction objects between spatial communicators.

For non-Firedrake objects these wrappers will dispatch to the normal implementations on mpi4py.MPI.Comm, which means that the same call site can be used for both Firedrake and non-Firedrake types.

Parameters:

  • comm – The communicator to split.

  • M – The size of the communicators used for spatial parallelism. Must be an integer divisor of the size of comm.

  • kwargs – Can include an ensemble_name string used as a communicator name prefix, for debugging.

Raises:

ValueError – If M does not divide comm.size exactly.

allreduce(f: Function | Cofunction, f_reduced: Function | Cofunction | None = None, op: Op = <mpi4py.MPI.Op object>) → Function | Cofunction[source]

Allreduce a Function f into f_reduced.

Parameters:

  • f – The Function to allreduce.

  • f_reduced – The result of the reduction. Must be in the same FunctionSpace() as f.

  • op – MPI reduction operator. Defaults to MPI.SUM.

Returns:

The result of the reduction.

Return type:

Function | Cofunction

Raises:

ValueError – If Function communicators mismatch each other or the ensemble spatial communicator, or if the Functions are in different spaces

bcast(f: Function | Cofunction, root: int = 0) → Function | Cofunction[source]

Broadcast a Function f over ensemble_comm from ensemble_rank root.

Parameters:

  • f – The Function to broadcast.

  • root – The rank to broadcast from.

Returns:

The result of the broadcast.

Return type:

Function | Cofunction

Raises:

ValueError – If the Function communicator mismatches the ensemble.comm.

property ensemble_rank: int

The rank of the local ensemble member.

property ensemble_size: int

The number of ensemble members.

iallreduce(f: Function | Cofunction, f_reduced: Function | Cofunction | None = None, op: Op = <mpi4py.MPI.Op object>) → list[Request][source]

Allreduce (non-blocking) a Function f into f_reduced.

Parameters:

  • f – The a Function to allreduce.

  • f_reduced – The result of the reduction. Must be in the same FunctionSpace() as f.

  • op – MPI reduction operator. Defaults to MPI.SUM.

Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If Function communicators mismatch each other or the ensemble spatial communicator, or if the Functions are in different spaces

ibcast(f: Function | Cofunction, root: int = 0) → list[Request][source]

Broadcast (non-blocking) a Function f over ensemble_comm ensemble_rank root.

Parameters:

  • f – The Function to broadcast.

  • root – The rank to broadcast from.

Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If the Function communicator mismatches the ensemble.comm.

irecv(f: Function | Cofunction, source: int = -1, tag: int = -1) → list[Request][source]

Receive (non-blocking) a Function f over ensemble_comm from another ensemble_rank.

Parameters:

  • f – The Function to receive into.

  • source – The ensemble_rank to receive f from.

  • tag – The tag of the message.

Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If the Function communicator mismatches the ensemble.comm.

ireduce(f: Function | Cofunction, f_reduced: Function | Cofunction | None = None, op: Op = <mpi4py.MPI.Op object>, root: int = 0) → list[Request][source]

Reduce (non-blocking) a Function f into f_reduced.

Parameters:

  • f – The a Function to reduce.

  • f_reduced – The result of the reduction. Must be in the same FunctionSpace() as f.

  • op – MPI reduction operator. Defaults to MPI.SUM.

  • root – The ensemble rank to reduce to.

Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If Function communicators mismatch each other or the ensemble spatial communicator, or if the Functions are in different spaces

isend(f: Function | Cofunction, dest: int, tag: int = 0) → list[Request][source]

Send (non-blocking) a Function f over ensemble_comm to another ensemble_rank.

Parameters:
Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If the Function communicator mismatches the ensemble.comm.

isendrecv(fsend: Function | Cofunction, dest: int, sendtag: int = 0, frecv: Function | Cofunction | None = None, source: int = -1, recvtag: int = -1) → list[Request][source]

Send (non-blocking) a Function fsend and receive a Function frecv over ensemble_comm to/from other ensemble_rank.

fsend and frecv do not need to be in the same function space.

Parameters:

  • fsend – The a Function to send.

  • dest – The ensemble_rank to send fsend to.

  • sendtag – The tag of the send message.

  • frecv – The Function to receive into.

  • source – The ensemble_rank to receive frecv from.

  • recvtag – The tag of the receive message.

Returns:

Requests one for each of f.subfunctions.

Return type:

list[mpi4py.MPI.Request]

Raises:

ValueError – If the Function communicators mismatches each other or the ensemble.comm.

recv(f: Function | Cofunction, source: int = -1, tag: int = -1, statuses: list[Status] | Status = None) → Function | Cofunction[source]

Receive (blocking) a Function f over ensemble_comm from another ensemble_rank.

Parameters:

  • f – The Function to receive into.

  • source – The ensemble_rank to receive f from.

  • tag – The tag of the message.

  • statuses – The mpi4py.MPI.Status of the internal recv calls (one for each of the subfunctions of f).

Returns:

f with the received data.

Return type:

Function | Cofunction

Raises:

  • ValueError – If the Function communicator mismatches the ensemble.comm.

  • ValueError – If the number of statuses provided is not the number of subfunctions of f.

reduce(f: Function | Cofunction, f_reduced: Function | Cofunction | None = None, op: Op = <mpi4py.MPI.Op object>, root: int = 0) → Function | Cofunction[source]

Reduce a Function f into f_reduced.

Parameters:

  • f – The Function to reduce.

  • f_reduced – The result of the reduction. Must be in the same FunctionSpace() as f.

  • op – MPI reduction operator. Defaults to MPI.SUM.

  • root – The ensemble rank to reduce to.

Returns:

The result of the reduction.

Return type:

Function | Cofunction

Raises:

ValueError – If Function communicators mismatch each other or the ensemble spatial communicator, or if the Functions are in different spaces

send(f: Function | Cofunction, dest: int, tag: int = 0)[source]

Send (blocking) a Function f over ensemble_comm to another ensemble_rank.

Parameters:
Raises:

ValueError – If the Function communicator mismatches the ensemble.comm.

sendrecv(fsend: Function | Cofunction, dest: int, sendtag: int = 0, frecv: Function | Cofunction | None = None, source: int = -1, recvtag: int = -1, statuses: list[Status] | Status = None) → Function | Cofunction[source]

Send (blocking) a Function fsend and receive a Function frecv over ensemble_comm to/from other ensemble_rank.

fsend and frecv do not need to be in the same function space but do need to have the same number of subfunctions.

Parameters:

  • fsend – The a Function to send.

  • dest – The ensemble_rank to send fsend to.

  • sendtag – The tag of the send message.

  • frecv – The Function to receive into.

  • source – The ensemble_rank to receive frecv from.

  • recvtag – The tag of the receive message.

  • statuses – The mpi4py.MPI.Status of the internal recv calls (one for each of the subfunctions of frecv).

Returns:

frecv with the received data.

Return type:

Function | Cofunction

Raises:

  • ValueError – If the Function communicators mismatches each other or the ensemble.comm.

  • ValueError – If the number of statuses provided is not the number of subfunctions of f.

sequential(*, synchronise: bool = False, reverse: bool = False, **kwargs)[source]

Context manager for executing code on each ensemble member consecutively (ordered by increasing ensemble_rank).

Any data in kwargs will be made available in the returned context and will be communicated forward after each ensemble member exits. Function or Cofunction kwargs will be sent with the corresponding Ensemble methods.

For example:

with ensemble.sequential(index=0) as ctx:
    print(ensemble.ensemble_rank, ctx.index)
    ctx.index += 2

Would print:

0 0
1 2
2 4
3 6
...

If reverse is True then the ensemble ranks will be looped through in decreasing order i.e. ensemble_rank == (ensemble_size - 1) will run first, then ensemble_rank == (ensemble_size - 2) etc.

Parameters:

  • synchronise – If True then MPI_Barrier will be called on the global_comm at the beginning and end of this method.

  • reverse – If True then will iterate through spatial comms in order of decreasing ensemble_rank.

  • kwargs – Data to be passed forward by each rank and made available in the returned ctx.

firedrake.ensemble.ensemble_function module

class firedrake.ensemble.ensemble_function.EnsembleCofunction(function_space: EnsembleDualSpace)[source]

Bases: EnsembleFunctionBase

A mixed finite element Cofunction distributed over an ensemble.

Parameters:

function_space (EnsembleDualSpace) – The function space of the cofunction.

class firedrake.ensemble.ensemble_function.EnsembleFunction(function_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleFunctionBase

A mixed Function defined on a Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally in an EnsembleFunctionSpace.

Parameters:

function_space (EnsembleFunctionSpace.) – The function space of the Function.

Notes

Passing an EnsembleDualSpace to EnsembleFunction will return an instance of EnsembleCofunction.

This class does not carry UFL symbolic information, unlike a Function. UFL expressions can only be defined locally on each ensemble member using a Function from EnsembleFunction.subfunctions.

See also

EnsembleFunctionSpace, EnsembleFunction, EnsembleDualSpace, EnsembleCofunction

norm(*args, **kwargs)[source]

Compute the norm of the function.

Any arguments are forwarded to norm().

class firedrake.ensemble.ensemble_function.EnsembleFunctionBase(function_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleFunctionMixin

A mixed (co)function defined on a Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally in an EnsembleFunctionSpace.

Parameters:

function_space – The function space of the (co)function.

Notes

Passing an EnsembleDualSpace to EnsembleFunction will return an instance of EnsembleCofunction.

This class does not carry UFL symbolic information, unlike a Function. UFL expressions can only be defined locally on each ensemble member using a \(~firedrake.function.Function\) from EnsembleFunction.subfunctions.

See also

ensemble_functionspace.EnsembleFunctionSpace, ensemble_functionspace.EnsembleDualSpace, EnsembleFunction, EnsembleCofunction

assign(other, subsets=None)[source]

Set the EnsembleFunction to the value of another EnsembleFunction other.

Parameters:
copy()[source]

Return a deep copy of the EnsembleFunction.

function_space()[source]
riesz_representation(**kwargs)[source]

Return the Riesz representation of this EnsembleFunction with respect to the given Riesz map.

Internally delegates to the firedrake.function.Function.riesz_representation() of each component.

Parameters:

  • riesz_map – The Riesz map to use (\(l2\), \(L2\), or \(H1\)). This can also be a callable.

  • kwargs – other arguments to be passed to the firedrake.riesz_map.

property subfunctions

The (co)functions on the local ensemble member.

vec()[source]

Context manager for the global PETSc.Vec with read/write access.

It is invalid to access the Vec outside of a context manager.

vec_ro()[source]

Context manager for the global PETSc.Vec with read only access.

It is invalid to access the Vec outside of a context manager.

vec_wo()[source]

Context manager for the global PETSc.Vec with write only access.

It is invalid to access the Vec outside of a context manager.

zero(subsets=None)[source]

Set values to zero.

Parameters:

subsets (Collection[Optional[pyop2.types.set.Subset]]) – One subset for each local firedrake.function.Function. None elements will be ignored. The values of each local function will only be zeroed on the nodes on the corresponding subset.

firedrake.ensemble.ensemble_functionspace module

class firedrake.ensemble.ensemble_functionspace.EnsembleDualSpace(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: EnsembleFunctionSpaceBase

A mixed dual function space defined on an ensemble.Ensemble. The subcomponents are distributed over the ensemble members, but are specified locally on each ensemble member.

Parameters:

  • local_spaces (Collection) – The list of dual function spaces on the local ensemble.comm.

  • ensemble (\(.ensemble.Ensemble\)) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FiredrakeDualSpace. UFL expressions can only be defined locally on each ensemble member using a FiredrakeDualSpace from \(EnsembleDualSpace.local_spaces\).

Examples

If U=CG1, V=DG0, and W=U*V, we can have the nested mixed dual space U*xV*xV*xW*xU*. This can be distributed over an ensemble.Ensemble with two ensemble members by splitting into [U*xV*]x[V*xW*xU*]. The following code creates the corresponding EnsembleDualSpace:

ensemble = Ensemble(COMM_WORLD, COMM_WORLD.size//2)
mesh = UnitIntervalMesh(8, comm=ensemble.comm)
U = FunctionSpace(mesh, "CG", 1)
V = FunctionSpace(mesh, "DG", 0)
W = U*V

if ensemble.ensemble_rank == 0:
    local_spaces = [U.dual(), V.dual()]
else:
    local_spaces = [V.dual(), W.dual(), U.dual()]

efs0 = EnsembleDualSpace(local_spaces, ensemble)
efs1 = EnsembleFunctionSpace(local_spaces, ensemble)

See also

EnsembleFunctionSpace, ensemble_function.EnsembleFunction, EnsembleDualSpace, ensemble_function.EnsembleCofunction

class firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpace(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: EnsembleFunctionSpaceBase

A mixed primal function space defined on an Ensemble. The subcomponents are distributed over the ensemble members, but are specified locally on each ensemble member.

Parameters:

  • local_spaces (Collection) – The list of primal function spaces on the local Ensemble.comm.

  • ensemble (Ensemble) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FunctionSpace. UFL expressions can only be defined locally on each ensemble member using a FunctionSpace from \(EnsembleFunctionSpace.local_spaces\).

Examples

If U=CG1, V=DG0, and W=UxV, we can have the nested mixed space UxVxVxWxU. This can be distributed over an ensemble.Ensemble with two ensemble members by splitting into [UxV]x[VxWxU]. The following code creates the corresponding EnsembleFunctionSpace:

ensemble = Ensemble(COMM_WORLD, COMM_WORLD.size//2)
mesh = UnitIntervalMesh(8, comm=ensemble.comm)
U = FunctionSpace(mesh, "CG", 1)
V = FunctionSpace(mesh, "DG", 0)
W = U*V

if ensemble.ensemble_rank == 0:
    local_spaces = [U, V]
else:
    local_spaces = [V, W, U]

efs = EnsembleFunctionSpace(local_spaces, ensemble)

See also

EnsembleFunctionSpace, ensemble_function.EnsembleFunction, EnsembleDualSpace, ensemble_function.EnsembleCofunction

class firedrake.ensemble.ensemble_functionspace.EnsembleFunctionSpaceBase(local_spaces: Collection, ensemble: Ensemble)[source]

Bases: object

Base class for mixed function spaces defined on an Ensemble. The subcomponents are distributed over the ensemble members, and are specified locally.

Parameters:

  • local_spaces (Collection) – The list of function spaces on the local ensemble.comm.

  • ensemble (\(~.ensemble.Ensemble\)) – The communicator that the function space is defined over.

Notes

Passing a list of dual local_spaces to EnsembleFunctionSpace will return an instance of EnsembleDualSpace.

This class does not carry UFL symbolic information, unlike a FunctionSpace. UFL expressions can only be defined locally on each ensemble member using a FunctionSpace from \(EnsembleFunctionSpace.local_spaces\).

See also

EnsembleFunctionSpace, EnsembleDualSpace, ensemble_function.EnsembleFunction, ensemble_function.EnsembleCofunction

property comm

The spatial communicator from the Ensemble communicator.

create_vec()[source]

Return a PETSc Vec on the Ensemble.global_comm with the same layout as an EnsembleFunction or EnsembleCofunction in this function space.

dual()[source]

The dual to this function space. A EnsembleDualSpace if self is a EnsembleFunctionSpace, and vice-versa.

property ensemble

The Ensemble that the function space is defined over

property ensemble_comm

The ensemble communicator from the Ensemble communicator.

property global_comm

The global communicator from the Ensemble communicator.

property global_spaces_offset

Index of the first local subspace in the global mixed space.

property layout_vec
property local_spaces

The FunctionSpace on the local ensemble.comm.

mesh()[source]

The Mesh on the local ensemble.comm.

property nglobal_dofs

The total number of dofs across all subspaces on all ensemble ranks.

property nglobal_spaces

The total number of subspaces across all ensemble ranks.

property nlocal_comm_dofs

The total number of dofs across all subspaces on the local ensemble.comm.

property nlocal_rank_dofs

The total number of dofs across all subspaces on the local MPI rank.

property nlocal_spaces

The number of subspaces on this ensemble rank.

firedrake.ensemble.ensemble_mat module

firedrake.ensemble.ensemble_mat.EnsembleBlockDiagonalMat(block_mats: Iterable, row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

A Mat for a block diagonal matrix defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace. This is a convenience function to create a PETSc.Mat with a EnsembleBlockDiagonalMatCtx Python context.

Parameters:

  • block_mats (Iterable[PETSc.Mat]) – The PETSc Mats for each block. On each ensemble rank there must be as many Mats as there are local subspaces of row_space and col_space, and the Mat sizes must match the sizes of the corresponding subspaces.

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Returns:

The PETSc.Mat with an EnsembleBlockDiagonalMatCtx Python context.

Return type:

PETSc.Mat

See also

EnsembleBlockDiagonalMatCtx, EnsembleBJacobiPC

class firedrake.ensemble.ensemble_mat.EnsembleBlockDiagonalMatCtx(block_mats: Iterable, row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

Bases: EnsembleMatCtxBase

A python Mat context for a block diagonal matrix defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace.

Parameters:

  • block_mats (Iterable[PETSc.Mat]) – The PETSc Mats for each block. On each ensemble rank there must be as many Mats as there are local subspaces of row_space and col_space, and the Mat sizes must match the sizes of the corresponding subspaces.

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Notes

This is a python context, not an actual PETSc.Mat. To create the corresponding PETSc.Mat users should call EnsembleBlockDiagonalMat().

See also

EnsembleBlockDiagonalMat, EnsembleBJacobiPC

mult_impl(A, x, y)[source]

Apply the action of the matrix to x, putting the result in y.

y is not guaranteed to be zero on entry. This is a convenience method allowing the matrix action to be implemented in terms of EnsembleFunction input and outputs by inheriting classes.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x – The vector acted on by the matrix.

  • y – The result of the matrix action.

See also

EnsembleMatCtxBase.mult

setUp(mat)[source]
view(mat, viewer=None)[source]
class firedrake.ensemble.ensemble_mat.EnsembleMatCtxBase(row_space: EnsembleFunctionSpaceBase, col_space: EnsembleFunctionSpaceBase)[source]

Bases: object

Base class for python type Mats defined over an Ensemble.

Parameters:

  • row_space – The function space that the matrix acts on. Must have the same number of subspaces on each ensemble rank as col_space.

  • col_space – The function space for the result of the matrix action. Must have the same number of subspaces on each ensemble rank as row_space.

Notes

The main use of this base class is to enable users to implement the matrix action as acting on and resulting in an EnsembleFunction. This is done by implementing the mult_impl method.

See also

ensemble_pc.EnsemblePCBase

mult(A, x, y)[source]

Apply the action of the matrix to x, putting the result in y.

This method will be called by PETSc with x and y as Vecs, and acts as a wrapper around the mult_impl method which has x and y as EnsembleFunction for convenience. y is not guaranteed to be zero on entry.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x (PETSc.Vec) – The vector acted on by the matrix.

  • y (PETSc.Vec) – The result of the matrix action.

See also

EnsembleMatCtxBase.mult_impl

mult_impl(A, x: EnsembleFunctionBase, y: EnsembleFunctionBase)[source]

Apply the action of the matrix to x, putting the result in y.

y is not guaranteed to be zero on entry. This is a convenience method allowing the matrix action to be implemented in terms of EnsembleFunction input and outputs by inheriting classes.

Parameters:

  • A (PETSc.Mat) – The PETSc matrix that self is the python context of.

  • x – The vector acted on by the matrix.

  • y – The result of the matrix action.

See also

EnsembleMatCtxBase.mult

firedrake.ensemble.ensemble_pc module

class firedrake.ensemble.ensemble_pc.EnsembleBJacobiPC[source]

Bases: EnsemblePCBase

A python PC context for a block Jacobi method defined over an Ensemble. Each block acts on a single subspace of an EnsembleFunctionSpace and is (approximately) solved with its own KSP, which defaults to -ksp_type preonly.

Available options:

  • -pc_use_amat - use Amat to apply block of operator in inner Krylov method

  • -sub_%d - set options for the %d’th block, numbered from ensemble rank 0.

  • -sub_ - set default options for all blocks.

Notes

Currently this is only implemented for EnsembleBlockDiagonalMatCtx matrices.

See also

EnsembleBlockDiagonalMatCtx, EnsembleBlockDiagonalMat

apply_impl(pc, x, y)[source]
initialize(pc)[source]

Initialize any state in this preconditioner.

This method is only called on the first time that the setUp method is called.

prefix = 'ebjacobi_'

The options prefix of this PC.

update(pc)[source]

Update any state in this preconditioner.

This method is called the on second and later times that the setUp method is called.

This method is not needed for all preconditioners and can often be a no-op.

view(pc, viewer=None)[source]

Write a basic description of this PC.

class firedrake.ensemble.ensemble_pc.EnsemblePCBase[source]

Bases: PCBase

Base class for python type PCs defined over an Ensemble.

The pc operators must be python Mats with EnsembleMatCtxBase.

Notes

The main use of this base class is to enable users to implement the preconditioner action as acting on and resulting in an EnsembleFunction. This is done by implementing the apply_impl method.

See also

EnsembleMatCtxBase

apply(pc, x, y)[source]

Apply the preconditioner to x, putting the result in y.

Both x and y are PETSc Vecs, y is not guaranteed to be zero on entry.

apply_impl(pc, x, y)[source]
initialize(pc)[source]

Initialize any state in this preconditioner.

This method is only called on the first time that the setUp method is called.

needs_python_pmat = True

Set this to False if the P matrix needs to be Python (matfree).

firedrake.ensemble.ensemble_pc.obj_name(obj)[source]

Module contents