firedrake.ensemble package

Submodules

firedrake.ensemble.ensemble module

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

Bases: object

Create a set of space and ensemble subcommunicators.

Parameters:
  • comm – The communicator to split.

  • M – the size of the communicators used for spatial parallelism.

Keyword Arguments:

ensemble_name – string used as communicator name prefix, for debugging.

Raises:

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

allreduce(f, f_reduced, op=<mpi4py.MPI.Op object>)[source]

Allreduce a function f into f_reduced over ensemble_comm .

Parameters:
  • f – The a Function to allreduce.

  • f_reduced – the result of the reduction.

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

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or if the functions are in different spaces

bcast(f, root=0)[source]

Broadcast a function f over ensemble_comm from rank root

Parameters:
  • f – The Function to broadcast.

  • root – rank to broadcast from. Defaults to 0.

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

property ensemble_rank

The rank of the local ensemble member.

property ensemble_size

The number of ensemble members.

iallreduce(f, f_reduced, op=<mpi4py.MPI.Op object>)[source]

Allreduce (non-blocking) a function f into f_reduced over ensemble_comm .

Parameters:
  • f – The a Function to allreduce.

  • f_reduced – the result of the reduction.

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

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or if the functions are in different spaces

ibcast(f, root=0)[source]

Broadcast (non-blocking) a function f over ensemble_comm from rank root

Parameters:
  • f – The Function to broadcast.

  • root – rank to broadcast from. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

irecv(f, source=-1, tag=-1)[source]

Receive (non-blocking) a function f over ensemble_comm from another ensemble rank.

Parameters:
  • f – The a Function to receive into

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • tag – the tag of the message. Defaults to MPI.ANY_TAG.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

ireduce(f, f_reduced, op=<mpi4py.MPI.Op object>, root=0)[source]

Reduce (non-blocking) a function f into f_reduced over ensemble_comm to rank root

Parameters:
  • f – The a Function to reduce.

  • f_reduced – the result of the reduction on rank root.

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

  • root – rank to reduce to. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or is the functions are in different spaces

isend(f, dest, tag=0)[source]

Send (non-blocking) a function f over ensemble_comm to another ensemble rank.

Parameters:
  • f – The a Function to send

  • dest – the rank to send to

  • tag – the tag of the message. Defaults to 0.

Returns:

list of MPI.Request objects (one for each of f.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

isendrecv(fsend, dest, sendtag=0, frecv=None, source=-1, recvtag=-1)[source]

Send a function fsend and receive a function frecv over ensemble_comm to another ensemble rank.

Parameters:
  • fsend – The a Function to send.

  • dest – the rank to send to.

  • sendtag – the tag of the send message. Defaults to 0.

  • frecv – The a Function to receive into.

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • recvtag – the tag of the received message. Defaults to MPI.ANY_TAG.

Returns:

list of MPI.Request objects (one for each of fsend.subfunctions and frecv.subfunctions).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

recv(f, source=-1, tag=-1, statuses=None)[source]

Receive (blocking) a function f over ensemble_comm from another ensemble rank.

Parameters:
  • f – The a Function to receive into

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • tag – the tag of the message. Defaults to MPI.ANY_TAG.

  • statuses – MPI.Status objects (one for each of f.subfunctions or None).

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

reduce(f, f_reduced, op=<mpi4py.MPI.Op object>, root=0)[source]

Reduce a function f into f_reduced over ensemble_comm to rank root

Parameters:
  • f – The a Function to reduce.

  • f_reduced – the result of the reduction on rank root.

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

  • root – rank to reduce to. Defaults to 0.

Raises:

ValueError – if function communicators mismatch each other or the ensemble spatial communicator, or is the functions are in different spaces

send(f, dest, tag=0)[source]

Send (blocking) a function f over ensemble_comm to another ensemble rank.

Parameters:
  • f – The a Function to send

  • dest – the rank to send to

  • tag – the tag of the message. Defaults to 0

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

sendrecv(fsend, dest, sendtag=0, frecv=None, source=-1, recvtag=-1, status=None)[source]

Send (blocking) a function fsend and receive a function frecv over ensemble_comm to another ensemble rank.

Parameters:
  • fsend – The a Function to send.

  • dest – the rank to send to.

  • sendtag – the tag of the send message. Defaults to 0.

  • frecv – The a Function to receive into.

  • source – the rank to receive from. Defaults to MPI.ANY_SOURCE.

  • recvtag – the tag of the received message. Defaults to MPI.ANY_TAG.

  • status – MPI.Status object or None.

Raises:

ValueError – if function communicator mismatches the ensemble spatial communicator.

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 a 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.

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

Compute the norm of the function.

Any arguments are forwarded to norm().

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)
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)

Module contents