From 10daf5ad82308e0f6f1314e2be2cbbfded1e7315 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sat, 25 Jul 2026 16:01:23 -0400 Subject: [PATCH 01/35] Add regional mask support to Field --- components/omega/src/infra/Field.cpp | 40 ++++++++++++++++++++++++++++ components/omega/src/infra/Field.h | 23 ++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/components/omega/src/infra/Field.cpp b/components/omega/src/infra/Field.cpp index ccf2d9d7984b..9dd92fe3a959 100644 --- a/components/omega/src/infra/Field.cpp +++ b/components/omega/src/infra/Field.cpp @@ -836,6 +836,46 @@ std::shared_ptr FieldGroup::getFieldFromGroup( return Group->getField(FieldName); } +//------------------------------------------------------------------------------ +// Regional mask functions +//------------------------------------------------------------------------------ + +//------------------------------------------------------------------------------ +// Set regional mask for this field +void Field::setRegionalMask(const Array1DI4 &Mask) { + // Validate input mask is allocated + if (Mask.data() == nullptr) { + LOG_ERROR("Field::setRegionalMask: attempting to set unallocated mask " + "for field {}", + FldName); + return; + } + // Shallow copy - just copy the view, not the data + RegionalMask = Mask; + HasRegionalMaskSet = true; +} + +//------------------------------------------------------------------------------ +// Get regional mask for this field +Array1DI4 Field::getRegionalMask() const { + // Defensive check: if mask was set but pointer is now null, that's a bug + if (HasRegionalMaskSet && RegionalMask.data() == nullptr) { + LOG_CRITICAL("Field::getRegionalMask: mask was set for field {} but is " + "now deallocated - lifetime bug!", + FldName); + ABORT("Field regional mask lifetime error"); + } + return RegionalMask; +} + +//------------------------------------------------------------------------------ +// Check if field has a regional mask +bool Field::hasRegionalMask() const { + // Return whether setRegionalMask() was called, not pointer state + // This way we can detect if mask *should* be there but isn't + return HasRegionalMaskSet; +} + //------------------------------------------------------------------------------ } // namespace OMEGA diff --git a/components/omega/src/infra/Field.h b/components/omega/src/infra/Field.h index 52a3b9a315a9..369ad183330c 100644 --- a/components/omega/src/infra/Field.h +++ b/components/omega/src/infra/Field.h @@ -89,6 +89,14 @@ class Field { /// various types and cast to the appropriate type when needed. std::shared_ptr DataArray; + /// Optional 1D regional mask (horizontal dimension) for analysis + /// Empty if no regional restriction applies + Array1DI4 RegionalMask; + + /// Flag to track whether setRegionalMask() was called + /// Used to detect lifetime bugs (mask set but pointer deallocated) + bool HasRegionalMaskSet = false; + /// Fills every element of InDataArray with the standard fill value for the /// array's element type, then records that value in the field metadata. /// Called automatically by attachData() so inactive entries are initialized @@ -241,6 +249,21 @@ class Field { void setOptionalRead(bool OptRead ///< [in] optional-read flag value ); + //--------------------------------------------------------------------------- + // Regional mask functions + //--------------------------------------------------------------------------- + /// Set regional mask for this field. Mask is 1D integer array over + /// horizontal dimension (0 = excluded, 1 = included) + void setRegionalMask(const Array1DI4 &Mask); + + /// Get regional mask for this field. Returns 1D integer regional mask array + /// (check hasRegionalMask() first) + Array1DI4 getRegionalMask() const; + + /// Check if field has a regional mask. Returns true if regional mask is set, + /// false otherwise + bool hasRegionalMask() const; + //--------------------------------------------------------------------------- // Metadata functions //--------------------------------------------------------------------------- From c9e7fa9b092787e192794b71debfe7015ec963a2 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sat, 25 Jul 2026 17:17:02 -0400 Subject: [PATCH 02/35] Add CoordinateBinningOp --- .../analysis/operators/CoordinateBinningOp.h | 292 ++++++++++++++++++ components/omega/src/infra/Field.cpp | 2 +- 2 files changed, 293 insertions(+), 1 deletion(-) create mode 100644 components/omega/src/analysis/operators/CoordinateBinningOp.h diff --git a/components/omega/src/analysis/operators/CoordinateBinningOp.h b/components/omega/src/analysis/operators/CoordinateBinningOp.h new file mode 100644 index 000000000000..ffdb4d0a3f8a --- /dev/null +++ b/components/omega/src/analysis/operators/CoordinateBinningOp.h @@ -0,0 +1,292 @@ +#ifndef OMEGA_COORDINATEBINNINGOP_H +#define OMEGA_COORDINATEBINNINGOP_H + +//===-- analysis/operators/CoordinateBinningOp.h ----------------*- C++ -*-===// +// +// +/// \file +/// \brief Defines the CoordinateBinningOp operator for coordinate binning +/// +/// CoordinateBinningOp assigns mesh cells (or edges/vertices) to bins +/// based on a coordinate field (e.g., latitude). This operator is particularly +/// useful for meridional overturning circulation (MOC) calculations, where +/// cells are binned by latitude for subsequent accumulation and integration. +/// +/// The operator takes a 1D coordinate field (e.g., latCell) and outputs a 1D +/// integer array containing the bin index for each cell. Bin boundaries are +/// either auto-computed from the mesh extent or specified via configuration. +/// The binning computation occurs once during initialization and the results +/// are cached, as bin assignments are constant over time. +/// +/// Configuration: +/// - NumBins: Number of bins (required) +/// - MinBin: Minimum coordinate value for binning (optional, auto-computed if +/// not specified) +/// - MaxBin: Maximum coordinate value for binning (optional, auto-computed if +/// not specified) +/// +/// Example usage in operator chain: +/// \code +/// LatCell_CoordinateBinning +/// \endcode +/// where CoordinateBinning assigns cells to latitude bins for MOC calculation. +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" +#include "Reductions.h" + +namespace OMEGA { + +/// CoordinateBinningOp assigns mesh entities to bins based on a +/// coordinate field. Bin boundaries are either auto-computed from mesh extent +/// or user-specified. The binning is computed once during initialization and +/// cached. Output is a 1D integer array of bin indices. +template class CoordinateBinningOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output is always I4 (integer bin indices) + using OutputArrayT = Array1DI4; + + /// Constructs a CoordinateBinningOp operator. Reads binning configuration + /// (number of bins, optional min/max bounds), creates output Field for bin + /// indices, allocates output data array, and registers the output Field. + /// The actual binning computation is deferred to initialize() when mesh + /// information is available. The output Field name is constructed as + /// InputName + "_BinIndex". + CoordinateBinningOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("CoordinateBinning") { + + // Store input field names + InputNames = UpstreamNames; + + // Read required NumBins parameter + Error Err = Options.get("NumBins", NumBins); + if (Err.isFail()) { + ABORT_ERROR("CoordinateBinningOp: Required parameter 'NumBins' not " + "found in configuration"); + } + + if (NumBins <= 0) { + ABORT_ERROR("CoordinateBinningOp: NumBins must be positive, got {}", + NumBins); + } + + // Read optional MinBin parameter + Err = Options.get("MinBin", MinBin); + if (Err.isFail()) { + AutoComputeMin = true; + } else { + AutoComputeMin = false; + } + + // Read optional MaxBin parameter + Err = Options.get("MaxBin", MaxBin); + if (Err.isFail()) { + AutoComputeMax = true; + } else { + AutoComputeMax = false; + } + + // Retrieve input Field to get dimensions + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Validate input is 1D + if (ArrayT::rank != 1) { + ABORT_ERROR( + "CoordinateBinningOp: Input must be 1D (coordinate array), " + "got rank {}", + static_cast(ArrayT::rank)); + } + + // Get dimension info + std::vector InputDimNames; + InputField->getDimNames(InputDimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = InputNames[0] + "_BinIndex"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Create output Field with same horizontal dimension as input + // but with I4 type for bin indices + auto OutputField = + Field::create(OutputNames[0], + "Bin index for " + InputNames[0], // Description + "1", // Units (dimensionless) + "", // Standard name + 0, // Min valid value + NumBins - 1, // Max valid value + 1, // Rank + InputDimNames // Dimension names + ); + + // Store array size + ArraySize = InputData.extent(0); + + // Allocate output data array (bin indices for each cell) + OutputData = OutputArrayT(OutputNames[0] + "_out", ArraySize); + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + // Mark binning as not yet computed + BinningComputed = false; + + } // end constructor + + /// Initializes the operator by computing bin boundaries and assigning cells + /// to bins. If MinBin/MaxBin are not specified, computes them from the + /// global min/max of the coordinate field. Bin assignments are cached and + /// reused in subsequent compute() calls. + void initialize(const MachEnv *Env, ///< [in] machine environment + const HorzMesh *Mesh, ///< [in] horizontal mesh + const VertCoord *VCoord, ///< [in] vertical coordinate + Config Options ///< [in] operator-specific options + ) override { + + // Call base class initialize + AnalysisOperator::initialize(Env, Mesh, VCoord, Options); + + // Retrieve input coordinate field + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Auto-compute MinBin and/or MaxBin if not specified + if (AutoComputeMin || AutoComputeMax) { + // Determine index space from field dimension name + std::vector InputDimNames; + InputField->getDimNames(InputDimNames); + std::string IndexSpaceName = InputDimNames[0]; + + I4 NOwned = 0; + if (IndexSpaceName == "NCells") { + NOwned = Mesh->NCellsOwned; + } else if (IndexSpaceName == "NEdges") { + NOwned = Mesh->NEdgesOwned; + } else if (IndexSpaceName == "NVertices") { + NOwned = Mesh->NVerticesOwned; + } else { + ABORT_ERROR("CoordinateBinningOp: Unknown index space {}", + IndexSpaceName); + } + + // Compute global min/max over owned entities + std::vector IndxRange = {0, NOwned - 1}; + + if (AutoComputeMin) { + MinBin = + static_cast(globalMinVal(InputData, Comm, &IndxRange)); + } + + if (AutoComputeMax) { + MaxBin = + static_cast(globalMaxVal(InputData, Comm, &IndxRange)); + } + } + + // Add small margins to avoid boundary issues + const Real Margin = (MaxBin - MinBin) * 1.0e-6; + MinBin -= Margin; + MaxBin += Margin; + + // Compute bin width + BinWidth = (MaxBin - MinBin) / NumBins; + + if (BinWidth <= 0) { + ABORT_ERROR("CoordinateBinningOp: Invalid bin width {}. MinBin={}, " + "MaxBin={}, NumBins={}", + BinWidth, MinBin, MaxBin, NumBins); + } + + // Compute bin assignments for all entities + computeBinning(InputData); + + BinningComputed = true; + + } // end initialize + + /// Computes bin assignments. Since binning is constant in time, this + /// simply validates that initialization has occurred. The actual binning + /// is done once in initialize() and cached. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + if (!BinningComputed) { + ABORT_ERROR("CoordinateBinningOp::compute called before initialize"); + } + + // Update cache validity markers + // (binning is constant, so no recomputation needed) + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Computes bin index for each entity based on coordinate value + void computeBinning(const ArrayT &CoordData) { + + auto LocalOutput = OutputData; + auto LocalMinBin = MinBin; + auto LocalBinWidth = BinWidth; + auto LocalNumBins = NumBins; + + parallelFor( + {ArraySize}, KOKKOS_LAMBDA(int i) { + Real CoordVal = static_cast(CoordData(i)); + + // Compute bin index: floor((coord - minBin) / binWidth) + I4 BinIdx = + static_cast((CoordVal - LocalMinBin) / LocalBinWidth); + + // Clamp to valid range [0, NumBins-1] + if (BinIdx < 0) { + BinIdx = 0; + } else if (BinIdx >= LocalNumBins) { + BinIdx = LocalNumBins - 1; + } + + LocalOutput(i) = BinIdx; + }); + + } // end computeBinning + + /// Output data array holding bin indices for each entity + OutputArrayT OutputData; + + /// Number of spatial bins + I4 NumBins; + + /// Minimum coordinate value for binning + Real MinBin; + + /// Maximum coordinate value for binning + Real MaxBin; + + /// Bin width (derived from MinBin, MaxBin, NumBins) + Real BinWidth; + + /// Whether to auto-compute MinBin from mesh + bool AutoComputeMin; + + /// Whether to auto-compute MaxBin from mesh + bool AutoComputeMax; + + /// Total size of the coordinate array + I4 ArraySize; + + /// Whether binning has been computed (once in initialize()) + bool BinningComputed; + +}; // end class CoordinateBinningOp + +} // end namespace OMEGA + +#endif diff --git a/components/omega/src/infra/Field.cpp b/components/omega/src/infra/Field.cpp index 9dd92fe3a959..f06a567392c1 100644 --- a/components/omega/src/infra/Field.cpp +++ b/components/omega/src/infra/Field.cpp @@ -851,7 +851,7 @@ void Field::setRegionalMask(const Array1DI4 &Mask) { return; } // Shallow copy - just copy the view, not the data - RegionalMask = Mask; + RegionalMask = Mask; HasRegionalMaskSet = true; } From f02df5d989b0273e95be9313138238b2b5a1f6e4 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sat, 25 Jul 2026 18:30:57 -0400 Subject: [PATCH 03/35] Add PseudoToGeometricOp --- .../analysis/operators/PseudoToGeometricOp.h | 453 ++++++++++++++++++ 1 file changed, 453 insertions(+) create mode 100644 components/omega/src/analysis/operators/PseudoToGeometricOp.h diff --git a/components/omega/src/analysis/operators/PseudoToGeometricOp.h b/components/omega/src/analysis/operators/PseudoToGeometricOp.h new file mode 100644 index 000000000000..bf7e4e46c1ea --- /dev/null +++ b/components/omega/src/analysis/operators/PseudoToGeometricOp.h @@ -0,0 +1,453 @@ +#ifndef OMEGA_PSEUDOTOGEOMETRICOP_H +#define OMEGA_PSEUDOTOGEOMETRICOP_H + +//===-- analysis/operators/PseudoToGeometricOp.h ----------------*- C++ -*-===// +// +/// \file +/// \brief Defines the PseudoToGeometricOp operator for coordinate conversion +/// +/// PseudoToGeometricOp converts fields from pseudo-height coordinates to +/// geometric coordinates in Omega's non-Boussinesq formulation. This operator +/// is essential for diagnostics and analysis that require physical (geometric) +/// quantities rather than the prognostic pseudo-coordinate variables. +/// +/// In Omega, the vertical coordinate is based on pseudo-height (proportional +/// to pressure), and prognostic variables include pseudo-thickness and +/// pseudo-velocity. The relationship between pseudo and geometric quantities +/// is: +/// +/// geometric_quantity = (RhoSw * SpecVol) * pseudo_quantity +/// +/// where RhoSw is the reference density and SpecVol is the in-situ specific +/// volume (SpecVol = 1/Rho). This is equivalent to (RhoSw/Rho) * +/// pseudo_quantity. +/// +/// Vertical Grid Staggering: +/// The operator automatically handles vertical grid staggering. If the pseudo +/// field is at layer interfaces (NVertLayers+1) while specific volume is at +/// layer midpoints (NVertLayers), the operator interpolates specific volume to +/// interfaces using: +/// - Top interface (k=0): uses top layer specific volume +/// - Interior interfaces: averages adjacent layer specific volumes +/// - Bottom interface: uses bottom layer specific volume +/// +/// The operator takes one input: +/// 1. A pseudo-coordinate field (e.g., VerticalPseudoVelocity, PseudoThickness) +/// +/// The in-situ specific volume field (SpecVol) is fetched via the Field +/// registry +/// +/// Example usage in operator chain: +/// \code +/// VerticalPseudoVelocity_PseudoToGeometric +/// \endcode +/// where PseudoToGeometric converts pseudo vertical velocity to geometric. +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" +#include "GlobalConstants.h" + +namespace OMEGA { + +/// PseudoToGeometricOp converts pseudo-coordinate fields to geometric +/// coordinates using the specific volume scaling relationship from Omega's +/// non-Boussinesq formulation. Supports 1D (horizontal), 2D (horizontal × +/// vertical), and 3D (structured horizontal × vertical) arrays. Works on any +/// horizontal index space (cells, edges, or vertices). The conversion formula +/// is: Output = (RhoSw * SpecVol) * Input. +template class PseudoToGeometricOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type - same as input array type + using OutputArrayT = ArrayT; + + /// Constructs a PseudoToGeometricOp operator. Retrieves input field + /// dimensions and metadata, creates output Field for the + /// geometric-coordinate quantity, allocates output data array, and registers + /// the output Field. The output Field name is constructed as InputName + + /// "_Geometric". + PseudoToGeometricOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("PseudoToGeometric") { + + // Store input field names + // UpstreamNames[0] = pseudo-coordinate field to convert + InputNames = UpstreamNames; + + if (InputNames.size() < 1) { + ABORT_ERROR("PseudoToGeometricOp: Requires 1 input (pseudo field)"); + } + + // Retrieve input pseudo-coordinate field to get dimensions + auto PseudoField = Field::get(InputNames[0]); + auto PseudoData = PseudoField->template getDataArray(); + + // Support 1D (horizontal), 2D (horizontal × vertical), and 3D arrays + const I4 InputRank = ArrayT::rank; + + // Store dimensions based on rank + if (InputRank == 1) { + // 1D: horizontal only (cells/edges/vertices), no vertical structure + NHorizDim = PseudoData.extent(0); + NVertSize = 1; // No vertical dimension + } else if (InputRank == 2) { + // 2D: NHorizDim × NVertSize + NHorizDim = PseudoData.extent(0); + NVertSize = PseudoData.extent(1); + } else if (InputRank == 3) { + // 3D: Dim0 × Dim1 × NVertSize (flatten horizontal grid) + NHorizDim = PseudoData.extent(1); + NVertSize = PseudoData.extent(2); + } else { + ABORT_ERROR("PseudoToGeometricOp: Unsupported rank {}. " + "Supports 1D, 2D, and 3D arrays", + InputRank); + } + + // Get dimension info from pseudo field + auto NDims = PseudoField->getNumDims(); + std::vector DimNames; + PseudoField->getDimNames(DimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = InputNames[0] + "_Geometric"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get input metadata + std::string PseudoDescr, PseudoUnits, PseudoStdName; + ScalarT PseudoValidMin, PseudoValidMax; + PseudoField->getMetadata("Description", PseudoDescr); + PseudoField->getMetadata("Units", PseudoUnits); + PseudoField->getMetadata("StdName", PseudoStdName); + PseudoField->getMetadata("ValidMin", PseudoValidMin); + PseudoField->getMetadata("ValidMax", PseudoValidMax); + + // Create output Field with updated metadata + auto OutputField = + Field::create(OutputNames[0], + "Geometric conversion of" + PseudoDescr, // Description + PseudoUnits, // Units + PseudoStdName, // Standard name + PseudoValidMin, // Min valid + PseudoValidMax, // Max valid + NDims, // Rank + DimNames // Dimension names + ); + + // Allocate output data array matching input layout + OutputData = OutputArrayT(OutputNames[0] + "_out", PseudoData.layout()); + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + } // end constructor + + /// Initializes the operator after all Fields exist. Determines the index + /// space (cells, edges, or vertices) from the pseudo field's horizontal + /// dimension name and stores the appropriate MinLayer/MaxLayer arrays + /// from VertCoord and connectivity arrays from Mesh for horizontal + /// interpolation of SpecVol. + void initialize(const MachEnv *Env, ///< [in] machine environment + const HorzMesh *Mesh, ///< [in] horizontal mesh + const VertCoord *VCoord, ///< [in] vertical coordinate + Config Options ///< [in] operator config + ) override { + + // Call base class initialize to store Mesh, VCoord, etc. + AnalysisOperator::initialize(Env, Mesh, VCoord, Options); + + // Get array rank for conditional logic + constexpr I4 InputRank = ArrayT::rank; + + // Determine index space from horizontal dimension + std::string IndexSpaceName; + if constexpr (InputRank == 1) { + // 1D: first (only) dimension + auto PseudoField = Field::get(InputNames[0]); + std::vector DimNames; + PseudoField->getDimNames(DimNames); + IndexSpaceName = DimNames[0]; + } else { + // 2D/3D: horizontal is 2nd to last dimension + auto PseudoField = Field::get(InputNames[0]); + std::vector DimNames; + PseudoField->getDimNames(DimNames); + IndexSpaceName = DimNames[InputRank - 2]; + } + + // Set index space type and connectivity arrays + if (IndexSpaceName == "NCells") { + IndexSpaceType = IndexSpace::Cell; + if constexpr (InputRank > 1) { + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + } + } else if (IndexSpaceName == "NEdges") { + IndexSpaceType = IndexSpace::Edge; + CellsOnEdge = Mesh->CellsOnEdge; + if constexpr (InputRank > 1) { + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + } + } else if (IndexSpaceName == "NVertices") { + IndexSpaceType = IndexSpace::Vertex; + CellsOnVertex = Mesh->CellsOnVertex; + VertexDegree = Mesh->VertexDegree; + if constexpr (InputRank > 1) { + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + } + } else { + ABORT_ERROR("PseudoToGeometricOp: Unknown index space {}", + IndexSpaceName); + } + + } // end initialize + + /// Computes the conversion from pseudo to geometric coordinates by + /// multiplying the pseudo field by (RhoSw * SpecVol) at each horizontal + /// point and vertical layer. Uses hierarchical parallelism with outer loop + /// over horizontal dimension and inner loop over vertical layers respecting + /// MinLayer/MaxLayer bounds. For fields at layer interfaces (like vertical + /// velocity), specific volume is interpolated from layer midpoints to + /// interfaces. Updates output data, timestamp, and computed flag. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Get array rank for conditional logic + constexpr I4 InputRank = ArrayT::rank; + + // Retrieve input pseudo field + auto PseudoField = Field::get(InputNames[0]); + auto PseudoData = PseudoField->template getDataArray(); + + // Retrieve SpecVol from Field registry + auto SpecVolField = Field::get("SpecVol"); + auto SpecVolData = SpecVolField->getDataArray(); + + // Check horizontal dimension compatibility + if (SpecVolData.extent(0) != NHorizDim) { + ABORT_ERROR( + "PseudoToGeometricOp: SpecVol field horizontal dimension " + "({}) does not match pseudo field horizontal dimension ({})", + SpecVolData.extent(0), NHorizDim); + } + + auto Output = OutputData; + + // Create functor for horizontal interpolation + HorizInterpFunctor getSpecVol{SpecVolData, IndexSpaceType, CellsOnEdge, + CellsOnVertex, VertexDegree}; + + if constexpr (InputRank == 1) { + // 1D case: horizontal field only, no vertical structure + parallelFor( + {NHorizDim}, KOKKOS_LAMBDA(int iHoriz) { + Real SpecVolHoriz = getSpecVol(iHoriz, 0); + Output(iHoriz) = + static_cast((RhoSw * SpecVolHoriz) * + static_cast(PseudoData(iHoriz))); + }); + + } else if constexpr (InputRank == 2) { + // 2D case: hierarchical parallelism over horizontal × vertical + + // Check vertical staggering + const I4 SpecVolNVertSize = SpecVolData.extent(1); + const bool AtInterfaces = (NVertSize == SpecVolNVertSize + 1); + + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + OMEGA_SCOPE(LocNVertSize, NVertSize); + + parallelForOuter( + "PseudoToGeometric2D", LaunchConfig({NHorizDim}), + KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { + const I4 KMin = LocMinLayer(iHoriz); + const I4 KMax = LocMaxLayer(iHoriz); + const I4 KRange = vertRange(KMin, KMax); + + parallelForInner( + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + + // Horizontal interpolation first, then vertical if + // needed + Real SpecVolAtK; + if (AtInterfaces) { + if (K == 0) { + // Top interface: horizontally interpolate at top + // layer + SpecVolAtK = getSpecVol(iHoriz, 0); + } else if (K == LocNVertSize - 1) { + // Bottom interface: horizontally interpolate at + // bottom + SpecVolAtK = + getSpecVol(iHoriz, SpecVolNVertSize - 1); + } else { + // Interior interface: horiz interp then vert + // average + Real SpecVolAbove = getSpecVol(iHoriz, K - 1); + Real SpecVolBelow = getSpecVol(iHoriz, K); + SpecVolAtK = + 0.5_Real * (SpecVolAbove + SpecVolBelow); + } + } else { + // Same vertical staggering: horizontal interpolation + // only + SpecVolAtK = getSpecVol(iHoriz, K); + } + + Output(iHoriz, K) = static_cast( + (RhoSw * SpecVolAtK) * + static_cast(PseudoData(iHoriz, K))); + }); + }); + + } else if constexpr (InputRank == 3) { + // 3D case: hierarchical parallelism over tracers/dim0 × horizontal + // Following the pattern from TimeStepper.cpp for tracer loops + + // Check vertical staggering + const I4 SpecVolNVertSize = SpecVolData.extent(1); + const bool AtInterfaces = (NVertSize == SpecVolNVertSize + 1); + + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + OMEGA_SCOPE(LocNVertSize, NVertSize); + + const I4 Dim0 = PseudoData.extent(0); + + parallelForOuter( + "PseudoToGeometric3D", LaunchConfig({Dim0, NHorizDim}), + KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { + const I4 KMin = LocMinLayer(iHoriz); + const I4 KMax = LocMaxLayer(iHoriz); + const I4 KRange = vertRange(KMin, KMax); + + parallelForInner( + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + + // Horizontal interpolation first, then vertical if + // needed + Real SpecVolAtK; + if (AtInterfaces) { + if (K == 0) { + // Top interface: horizontally interpolate at top + // layer + SpecVolAtK = getSpecVol(iHoriz, 0); + } else if (K == LocNVertSize - 1) { + // Bottom interface: horizontally interpolate at + // bottom + SpecVolAtK = + getSpecVol(iHoriz, SpecVolNVertSize - 1); + } else { + // Interior interface: horiz interp then vert + // average + Real SpecVolAbove = getSpecVol(iHoriz, K - 1); + Real SpecVolBelow = getSpecVol(iHoriz, K); + SpecVolAtK = + 0.5_Real * (SpecVolAbove + SpecVolBelow); + } + } else { + // Same vertical staggering: horizontal interpolation + // only + SpecVolAtK = getSpecVol(iHoriz, K); + } + + Output(i0, iHoriz, K) = static_cast( + (RhoSw * SpecVolAtK) * + static_cast(PseudoData(i0, iHoriz, K))); + }); + }); + } + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Enum for horizontal index space type + enum class IndexSpace { Cell, Edge, Vertex }; + + /// This functor holds all necessary data to perform horizontal + /// interpolation of SpecVol. + struct HorizInterpFunctor { + Array2DReal SpecVolData; + IndexSpace IndexSpaceType; + Array2DI4 CellsOnEdge; + Array2DI4 CellsOnVertex; + I4 VertexDegree; + + /// Horizontally interpolates SpecVol to the field's index space. + /// For cell-based fields, returns SpecVol directly. + /// For edge-based fields, averages from 2 adjacent cells. + /// For vertex-based fields, averages from VertexDegree adjacent cells. + KOKKOS_INLINE_FUNCTION + Real operator()(const I4 iHoriz, ///< Horizontal index + const I4 K ///< Vertical layer + ) const { + + if (IndexSpaceType == IndexSpace::Cell) { + // Cell-based: use directly + return SpecVolData(iHoriz, K); + + } else if (IndexSpaceType == IndexSpace::Edge) { + // Edge-based: average from 2 adjacent cells + const I4 Cell1 = CellsOnEdge(iHoriz, 0); + const I4 Cell2 = CellsOnEdge(iHoriz, 1); + return 0.5_Real * (SpecVolData(Cell1, K) + SpecVolData(Cell2, K)); + + } else { // IndexSpace::Vertex + // Vertex-based: average from VertexDegree adjacent cells + Real SpecVolSum = 0.0_Real; + for (int j = 0; j < VertexDegree; ++j) { + const I4 Cell = CellsOnVertex(iHoriz, j); + SpecVolSum += SpecVolData(Cell, K); + } + return SpecVolSum / static_cast(VertexDegree); + } + } + }; + + /// Output data array holding the geometric-coordinate field + OutputArrayT OutputData; + + /// Number of points in horizontal dimension (cells, edges, or vertices) + I4 NHorizDim; + + /// Vertical size of the pseudo field array + I4 NVertSize; + + /// Type of horizontal index space (cell, edge, or vertex) + IndexSpace IndexSpaceType; + + /// Connectivity: cells adjacent to each edge (for edge-based fields) + Array2DI4 CellsOnEdge; + + /// Connectivity: cells sharing each vertex (for vertex-based fields) + Array2DI4 CellsOnVertex; + + /// Number of cells sharing each vertex (for vertex-based fields) + I4 VertexDegree; + + /// Min active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MinLayer; + + /// Max active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MaxLayer; + +}; // end class PseudoToGeometricOp + +} // end namespace OMEGA + +#endif From 8d187aba73b29a024884e725d4599f2d02ea040d Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sat, 25 Jul 2026 22:49:49 -0400 Subject: [PATCH 04/35] Add BinaryMultiplyOp --- .../src/analysis/operators/BinaryMultiplyOp.h | 322 ++++++++++++++++++ components/omega/src/infra/Field.cpp | 7 +- 2 files changed, 325 insertions(+), 4 deletions(-) create mode 100644 components/omega/src/analysis/operators/BinaryMultiplyOp.h diff --git a/components/omega/src/analysis/operators/BinaryMultiplyOp.h b/components/omega/src/analysis/operators/BinaryMultiplyOp.h new file mode 100644 index 000000000000..85dec17c17eb --- /dev/null +++ b/components/omega/src/analysis/operators/BinaryMultiplyOp.h @@ -0,0 +1,322 @@ +#ifndef OMEGA_BINARYMULTIPLYOP_H +#define OMEGA_BINARYMULTIPLYOP_H + +//===-- analysis/operators/BinaryMultiplyOp.h -------------------*- C++ -*-===// +// +/// \file +/// \brief Defines the BinaryMultiplyOp operator for element-wise field +/// multiplication +/// +/// BinaryMultiplyOp performs element-wise multiplication of two fields with +/// matching dimensions. This operator is useful for computing products like: +/// - Flux calculations (velocity × area) +/// - Density-weighted quantities (field × density) +/// - Coordinate transformations (field1 × field2) +/// +/// The operator supports two modes: +/// 1. Same-rank multiplication: Both fields have the same dimensions +/// Output(i,j,...) = Field1(i,j,...) × Field2(i,j,...) +/// +/// 2. Vertical expansion: Field1 is 2D/3D and Field2 is 1D (horizontal only) +/// The 1D field value is replicated across all vertical layers +/// Output(i,j) = Field1(i,j) × Field2(i) +/// +/// Example usage in operator chain: +/// \code +/// VerticalVelocity_BinaryMultiply(AreaCell) +/// \endcode +/// Computes the vertical flux by multiplying 2D vertical velocity by 1D cell +/// area, with the area value replicated across all vertical layers. +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" + +namespace OMEGA { + +/// BinaryMultiplyOp performs element-wise multiplication of two fields. +/// Supports both same-rank multiplication and vertical expansion where 1D +/// fields are replicated across all vertical layers when multiplying with 2D/3D +/// fields. The output field has the same dimensions as the first input field. +template class BinaryMultiplyOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type - same as input array types + using OutputArrayT = ArrayT; + + /// Constructs a BinaryMultiplyOp operator. Retrieves both input field + /// dimensions and metadata, validates that they match, creates output + /// Field for the product, allocates output data array, and registers + /// the output Field. The output Field name is constructed as + /// Field1Name + "_BinaryMultiply(" + Field2Name + ")_Product". + BinaryMultiplyOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("BinaryMultiply") { + + // Store input field names + // UpstreamNames[0] = first field + // UpstreamNames[1] = second field + InputNames = UpstreamNames; + + if (InputNames.size() < 2) { + ABORT_ERROR("BinaryMultiplyOp: Requires 2 input fields"); + } + + // Retrieve both input fields + auto Field1 = Field::get(InputNames[0]); + auto Field2 = Field::get(InputNames[1]); + auto Data1 = Field1->template getDataArray(); + + // Get rank information for both fields + const I4 Rank1 = ArrayT::rank; + const I4 Rank2 = Field2->getNumDims(); + + // Store whether Field2 is 1D (will be expanded vertically if Field1 is + // 2D/3D) + IsVerticalExpansion = (Rank1 > 1 && Rank2 == 1); + + // Validate dimension compatibility + if (!IsVerticalExpansion && Rank1 != Rank2) { + ABORT_ERROR("BinaryMultiplyOp: Input fields must have same rank or " + "second field must be 1D for vertical expansion, got " + "ranks {} and {}", + Rank1, Rank2); + } + + // For same-rank case, validate all dimensions match + if (!IsVerticalExpansion) { + auto Data2 = Field2->template getDataArray(); + for (I4 dim = 0; dim < Rank1; ++dim) { + if (Data1.extent(dim) != Data2.extent(dim)) { + ABORT_ERROR("BinaryMultiplyOp: Dimension {} mismatch: " + "Field1={}, Field2={}", + dim, Data1.extent(dim), Data2.extent(dim)); + } + } + } else { + // For vertical expansion case, validate horizontal dimensions match + auto Data2_1D = Field2->template getDataArray(); + if (Data1.extent(Rank1 - 2) != Data2_1D.extent(0)) { + ABORT_ERROR("BinaryMultiplyOp: Horizontal dimension mismatch for " + "vertical expansion: " + "Field1 horizontal dim={}, Field2 size={}", + Data1.extent(Rank1 - 2), Data2_1D.extent(0)); + } + } + + // Store dimensions based on rank + if (Rank1 == 1) { + // 1D: horizontal only (cells/edges/vertices), no vertical structure + NHorizDim = Data1.extent(0); + NVertSize = 1; // No vertical dimension + } else if (Rank1 == 2) { + // 2D: NHorizDim × NVertSize + NHorizDim = Data1.extent(0); + NVertSize = Data1.extent(1); + } else if (Rank1 == 3) { + // 3D: OuterDim × NHorizDim × NVertSize + NHorizDim = Data1.extent(1); + NVertSize = Data1.extent(2); + } else { + ABORT_ERROR("BinaryMultiplyOp: Unsupported rank {}. " + "Supports 1D, 2D, and 3D arrays", + Rank1); + } + + // Get dimension info from first field + auto NDims = Field1->getNumDims(); + std::vector DimNames; + Field1->getDimNames(DimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = + InputNames[0] + "_" + InputNames[1] + "_Product"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get metadata from both input fields + std::string Descr1, Units1, StdName1; + std::string Descr2, Units2, StdName2; + + Field1->getMetadata("Description", Descr1); + Field1->getMetadata("Units", Units1); + Field1->getMetadata("StdName", StdName1); + + Field2->getMetadata("Description", Descr2); + Field2->getMetadata("Units", Units2); + Field2->getMetadata("StdName", StdName2); + + // Combine metadata for output + std::string OutputDescr = "Product of " + Descr1 + " and " + Descr2; + std::string OutputUnits = combineUnits(Units1, Units2); + std::string OutputStdName = ""; // No standard name for generic product + + // Create output Field + auto OutputField = + Field::create(OutputNames[0], + OutputDescr, // Description + OutputUnits, // Units (combined) + OutputStdName, // Standard name + -std::numeric_limits::max(), // Min valid value + std::numeric_limits::max(), // Max valid value + NDims, // Rank + DimNames // Dimension names + ); + + // Allocate output data array matching input layout + OutputData = OutputArrayT(OutputNames[0] + "_out", Data1.layout()); + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + // Propagate regional mask from first input to output if present + // (Binary operations preserve the spatial structure of the first input) + if (Field1->hasRegionalMask()) { + OutputField->setRegionalMask(Field1->getRegionalMask()); + } + + } // end constructor + + /// Computes the element-wise multiplication by retrieving both input + /// data arrays and performing parallel multiplication using hierarchical + /// parallelism (outer loop over horizontal dimension, inner loop over + /// vertical layers). Supports 1D (horizontal only), 2D (horizontal × + /// vertical), and 3D (additional outer dimension × horizontal × vertical) + /// arrays. Also supports vertical expansion where a 1D field value is + /// replicated across all vertical layers when multiplying with 2D/3D fields. + /// Updates output data, timestamp, and computed flag. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Get array rank for conditional logic + constexpr I4 InputRank = ArrayT::rank; + + // Retrieve input Fields + auto Field1 = Field::get(InputNames[0]); + auto Field2 = Field::get(InputNames[1]); + auto Data1 = Field1->template getDataArray(); + + auto Output = OutputData; + + if constexpr (InputRank == 1) { + // 1D case: horizontal field only, no vertical structure + auto Data2 = Field2->template getDataArray(); + parallelFor( + {NHorizDim}, KOKKOS_LAMBDA(int iHoriz) { + Output(iHoriz) = + static_cast(static_cast(Data1(iHoriz)) * + static_cast(Data2(iHoriz))); + }); + + } else if constexpr (InputRank == 2) { + // 2D case: hierarchical parallelism over horizontal × vertical + + OMEGA_SCOPE(LocNVertSize, NVertSize); + + if (IsVerticalExpansion) { + // Vertical expansion: Data2 is 1D, replicate across vertical layers + auto Data2_1D = Field2->template getDataArray(); + parallelForOuter( + "BinaryMultiply2D_VertExpand", LaunchConfig({NHorizDim}), + KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { + const Real Data2Val = static_cast(Data2_1D(iHoriz)); + parallelForInner( + Team, LocNVertSize, INNER_LAMBDA(int K) { + Output(iHoriz, K) = static_cast( + static_cast(Data1(iHoriz, K)) * Data2Val); + }); + }); + } else { + // Same rank: element-wise multiplication + auto Data2 = Field2->template getDataArray(); + parallelForOuter( + "BinaryMultiply2D", LaunchConfig({NHorizDim}), + KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { + parallelForInner( + Team, LocNVertSize, INNER_LAMBDA(int K) { + Output(iHoriz, K) = static_cast( + static_cast(Data1(iHoriz, K)) * + static_cast(Data2(iHoriz, K))); + }); + }); + } + + } else if constexpr (InputRank == 3) { + // 3D case: hierarchical parallelism over dim0 × horizontal × vertical + + OMEGA_SCOPE(LocNVertSize, NVertSize); + const I4 Dim0 = Data1.extent(0); + + if (IsVerticalExpansion) { + // Vertical expansion: Data2 is 1D, replicate across vertical layers + auto Data2_1D = Field2->template getDataArray(); + parallelForOuter( + "BinaryMultiply3D_VertExpand", LaunchConfig({Dim0, NHorizDim}), + KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { + const Real Data2Val = static_cast(Data2_1D(iHoriz)); + parallelForInner( + Team, LocNVertSize, INNER_LAMBDA(int K) { + Output(i0, iHoriz, K) = static_cast( + static_cast(Data1(i0, iHoriz, K)) * + Data2Val); + }); + }); + } else { + // Same rank: element-wise multiplication + auto Data2 = Field2->template getDataArray(); + parallelForOuter( + "BinaryMultiply3D", LaunchConfig({Dim0, NHorizDim}), + KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { + parallelForInner( + Team, LocNVertSize, INNER_LAMBDA(int K) { + Output(i0, iHoriz, K) = static_cast( + static_cast(Data1(i0, iHoriz, K)) * + static_cast(Data2(i0, iHoriz, K))); + }); + }); + } + } + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Combines units from two fields for the product + std::string combineUnits(const std::string &Units1, + const std::string &Units2) { + if (Units1.empty() && Units2.empty()) { + return ""; + } else if (Units1.empty()) { + return Units2; + } else if (Units2.empty()) { + return Units1; + } else { + // Combine with multiplication notation + return Units1 + " * " + Units2; + } + } + + /// Output data array holding the product field values + OutputArrayT OutputData; + + /// Number of points in horizontal dimension (cells, edges, or vertices) + I4 NHorizDim; + + /// Vertical size of the array + I4 NVertSize; + + /// Whether Field2 is 1D and should be replicated across vertical dimension + bool IsVerticalExpansion; + +}; // end class BinaryMultiplyOp + +} // end namespace OMEGA + +#endif diff --git a/components/omega/src/infra/Field.cpp b/components/omega/src/infra/Field.cpp index f06a567392c1..6b000cd7697d 100644 --- a/components/omega/src/infra/Field.cpp +++ b/components/omega/src/infra/Field.cpp @@ -860,10 +860,9 @@ void Field::setRegionalMask(const Array1DI4 &Mask) { Array1DI4 Field::getRegionalMask() const { // Defensive check: if mask was set but pointer is now null, that's a bug if (HasRegionalMaskSet && RegionalMask.data() == nullptr) { - LOG_CRITICAL("Field::getRegionalMask: mask was set for field {} but is " - "now deallocated - lifetime bug!", - FldName); - ABORT("Field regional mask lifetime error"); + ABORT_ERROR("Field::getRegionalMask: mask was set for field {} but is " + "now deallocated.", + FldName); } return RegionalMask; } From 11602aa80077c8f64c751ab18c13d28e923b72e3 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sun, 26 Jul 2026 15:31:55 -0400 Subject: [PATCH 05/35] Add BinnedAccumulatorOp --- .../src/analysis/operators/BinaryMultiplyOp.h | 122 ++++-- .../analysis/operators/BinnedAccumulatorOp.h | 387 ++++++++++++++++++ 2 files changed, 478 insertions(+), 31 deletions(-) create mode 100644 components/omega/src/analysis/operators/BinnedAccumulatorOp.h diff --git a/components/omega/src/analysis/operators/BinaryMultiplyOp.h b/components/omega/src/analysis/operators/BinaryMultiplyOp.h index 85dec17c17eb..32c5677ff016 100644 --- a/components/omega/src/analysis/operators/BinaryMultiplyOp.h +++ b/components/omega/src/analysis/operators/BinaryMultiplyOp.h @@ -181,14 +181,48 @@ template class BinaryMultiplyOp : public AnalysisOperator { } // end constructor + /// Initializes the operator after all Fields exist. Determines the index + /// space (cells, edges, or vertices) from the first input field's horizontal + /// dimension name and stores the appropriate MinLayer/MaxLayer arrays + /// from VertCoord for bounding the inner vertical loop. + void initialize(const MachEnv *Env, const HorzMesh *InMesh, + const VertCoord *InVCoord, Config Options) override { + + AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); + + constexpr I4 InputRank = ArrayT::rank; + if constexpr (InputRank > 1) { + auto Field1 = Field::get(InputNames[0]); + std::vector DimNames; + Field1->getDimNames(DimNames); + // Horizontal dimension is 2nd-to-last for 2D/3D + std::string IndexSpaceName = DimNames[InputRank - 2]; + + if (IndexSpaceName == "NCells") { + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + } else if (IndexSpaceName == "NEdges") { + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + } else if (IndexSpaceName == "NVertices") { + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + } else { + ABORT_ERROR("BinaryMultiplyOp: Unknown index space {}", + IndexSpaceName); + } + } + + } // end initialize + /// Computes the element-wise multiplication by retrieving both input /// data arrays and performing parallel multiplication using hierarchical /// parallelism (outer loop over horizontal dimension, inner loop over - /// vertical layers). Supports 1D (horizontal only), 2D (horizontal × - /// vertical), and 3D (additional outer dimension × horizontal × vertical) - /// arrays. Also supports vertical expansion where a 1D field value is - /// replicated across all vertical layers when multiplying with 2D/3D fields. - /// Updates output data, timestamp, and computed flag. + /// vertical layers bounded by MinLayer/MaxLayer). Supports 1D (horizontal + /// only), 2D (horizontal × vertical), and 3D (additional outer dimension × + /// horizontal × vertical) arrays. Also supports vertical expansion where a + /// 1D field value is replicated across all vertical layers when multiplying + /// with 2D/3D fields. Updates output data, timestamp, and computed flag. void compute(const TimeInstant &TimeStamp ///< [in] current timestamp ) override { @@ -206,28 +240,34 @@ template class BinaryMultiplyOp : public AnalysisOperator { // 1D case: horizontal field only, no vertical structure auto Data2 = Field2->template getDataArray(); parallelFor( - {NHorizDim}, KOKKOS_LAMBDA(int iHoriz) { - Output(iHoriz) = - static_cast(static_cast(Data1(iHoriz)) * - static_cast(Data2(iHoriz))); + {NHorizDim}, KOKKOS_LAMBDA(int IHoriz) { + Output(IHoriz) = + static_cast(static_cast(Data1(IHoriz)) * + static_cast(Data2(IHoriz))); }); } else if constexpr (InputRank == 2) { // 2D case: hierarchical parallelism over horizontal × vertical + // Inner loop bounded by MinLayer/MaxLayer for partial columns - OMEGA_SCOPE(LocNVertSize, NVertSize); + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); if (IsVerticalExpansion) { // Vertical expansion: Data2 is 1D, replicate across vertical layers auto Data2_1D = Field2->template getDataArray(); parallelForOuter( "BinaryMultiply2D_VertExpand", LaunchConfig({NHorizDim}), - KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { - const Real Data2Val = static_cast(Data2_1D(iHoriz)); + KOKKOS_LAMBDA(int IHoriz, const TeamMember &Team) { + const Real Data2Val = static_cast(Data2_1D(IHoriz)); + const I4 KMin = LocMinLayer(IHoriz); + const I4 KMax = LocMaxLayer(IHoriz); + const I4 KRange = vertRange(KMin, KMax); parallelForInner( - Team, LocNVertSize, INNER_LAMBDA(int K) { - Output(iHoriz, K) = static_cast( - static_cast(Data1(iHoriz, K)) * Data2Val); + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + Output(IHoriz, K) = static_cast( + static_cast(Data1(IHoriz, K)) * Data2Val); }); }); } else { @@ -235,20 +275,26 @@ template class BinaryMultiplyOp : public AnalysisOperator { auto Data2 = Field2->template getDataArray(); parallelForOuter( "BinaryMultiply2D", LaunchConfig({NHorizDim}), - KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { + KOKKOS_LAMBDA(int IHoriz, const TeamMember &Team) { + const I4 KMin = LocMinLayer(IHoriz); + const I4 KMax = LocMaxLayer(IHoriz); + const I4 KRange = vertRange(KMin, KMax); parallelForInner( - Team, LocNVertSize, INNER_LAMBDA(int K) { - Output(iHoriz, K) = static_cast( - static_cast(Data1(iHoriz, K)) * - static_cast(Data2(iHoriz, K))); + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + Output(IHoriz, K) = static_cast( + static_cast(Data1(IHoriz, K)) * + static_cast(Data2(IHoriz, K))); }); }); } } else if constexpr (InputRank == 3) { // 3D case: hierarchical parallelism over dim0 × horizontal × vertical + // Inner loop bounded by MinLayer/MaxLayer for partial columns - OMEGA_SCOPE(LocNVertSize, NVertSize); + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); const I4 Dim0 = Data1.extent(0); if (IsVerticalExpansion) { @@ -256,12 +302,16 @@ template class BinaryMultiplyOp : public AnalysisOperator { auto Data2_1D = Field2->template getDataArray(); parallelForOuter( "BinaryMultiply3D_VertExpand", LaunchConfig({Dim0, NHorizDim}), - KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { - const Real Data2Val = static_cast(Data2_1D(iHoriz)); + KOKKOS_LAMBDA(int I0, int IHoriz, const TeamMember &Team) { + const Real Data2Val = static_cast(Data2_1D(IHoriz)); + const I4 KMin = LocMinLayer(IHoriz); + const I4 KMax = LocMaxLayer(IHoriz); + const I4 KRange = vertRange(KMin, KMax); parallelForInner( - Team, LocNVertSize, INNER_LAMBDA(int K) { - Output(i0, iHoriz, K) = static_cast( - static_cast(Data1(i0, iHoriz, K)) * + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + Output(I0, IHoriz, K) = static_cast( + static_cast(Data1(I0, IHoriz, K)) * Data2Val); }); }); @@ -270,12 +320,16 @@ template class BinaryMultiplyOp : public AnalysisOperator { auto Data2 = Field2->template getDataArray(); parallelForOuter( "BinaryMultiply3D", LaunchConfig({Dim0, NHorizDim}), - KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { + KOKKOS_LAMBDA(int I0, int IHoriz, const TeamMember &Team) { + const I4 KMin = LocMinLayer(IHoriz); + const I4 KMax = LocMaxLayer(IHoriz); + const I4 KRange = vertRange(KMin, KMax); parallelForInner( - Team, LocNVertSize, INNER_LAMBDA(int K) { - Output(i0, iHoriz, K) = static_cast( - static_cast(Data1(i0, iHoriz, K)) * - static_cast(Data2(i0, iHoriz, K))); + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + Output(I0, IHoriz, K) = static_cast( + static_cast(Data1(I0, IHoriz, K)) * + static_cast(Data2(I0, IHoriz, K))); }); }); } @@ -315,6 +369,12 @@ template class BinaryMultiplyOp : public AnalysisOperator { /// Whether Field2 is 1D and should be replicated across vertical dimension bool IsVerticalExpansion; + /// Min active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MinLayer; + + /// Max active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MaxLayer; + }; // end class BinaryMultiplyOp } // end namespace OMEGA diff --git a/components/omega/src/analysis/operators/BinnedAccumulatorOp.h b/components/omega/src/analysis/operators/BinnedAccumulatorOp.h new file mode 100644 index 000000000000..8e11460118b9 --- /dev/null +++ b/components/omega/src/analysis/operators/BinnedAccumulatorOp.h @@ -0,0 +1,387 @@ +#ifndef OMEGA_BINNEDACCUMULATOROP_H +#define OMEGA_BINNEDACCUMULATOROP_H + +//===-- analysis/operators/BinnedAccumulatorOp.h - BinnedAccumulatorOp -*- C++ -*-===// +// +/// \file +/// \brief Defines the BinnedAccumulatorOp operator for binned accumulation +/// +/// BinnedAccumulatorOp accumulates field values into spatial bins. This +/// operator is essential for MOC (Meridional Overturning Circulation) +/// calculations, where it accumulates edge normal velocities into latitude +/// bins. The operator performs local accumulation followed by MPI reduction +/// across all ranks to produce global bin totals. +/// +/// The operator takes two inputs: +/// 1. A value field to accumulate (1D or 2D) +/// 2. A bin index field (from CoordinateBinningOp) indicating which bin each +/// entity belongs to +/// +/// If the value field has a Field-attached regional mask it is applied +/// automatically during accumulation (1.0 = include, 0.0 = exclude). +/// +/// The output replaces the horizontal dimension with NumBins. For 2D input +/// (NCells x NVertLevels) the output is 2D (NumBins x NVertLevels); for 1D +/// input the output is 1D (NumBins). +/// +/// Configuration: +/// - NumBins: Number of spatial bins (required) +/// +/// Example usage in operator chain: +/// \code +/// NormalVelocity_BinnedAccumulator +/// \endcode +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" +#include "mpi.h" + +namespace OMEGA { + +/// BinnedAccumulatorOp accumulates field values into spatial bins. Supports +/// 1D and 2D input fields and produces output with the horizontal dimension +/// replaced by NumBins. If the input Field has a Field-attached regional mask +/// it is applied automatically. Uses hierarchical parallelism for 2D input. +/// Performs MPI reduction to compute global bin totals. This is the core +/// operator for MOC streamfunction calculations. +template class BinnedAccumulatorOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type depends on input rank: + /// For 2D input (e.g., NCells x NVertLevels), output is 2D (NumBins x NVertLevels) + /// For 1D input (e.g., NCells), output is 1D (NumBins) + using OutputArrayT = std::conditional_t; + + /// Constructs a BinnedAccumulatorOp operator. Reads the required NumBins + /// configuration parameter, creates the output Field for binned + /// accumulation, allocates output and local accumulation arrays, and + /// registers the output Field. The output Field name is constructed as + /// InputName + "_BinnedAccum". + BinnedAccumulatorOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("BinnedAccumulator") { + + // Store input field names: + // UpstreamNames[0] = value field to accumulate + // UpstreamNames[1] = bin index field (from CoordinateBinningOp) + InputNames = UpstreamNames; + + if (InputNames.size() < 2) { + ABORT_ERROR("BinnedAccumulatorOp: Requires at least 2 inputs (value " + "field and bin index field)"); + } + + // Read required NumBins parameter + Error Err = Options.get("NumBins", NumBins); + if (Err.isFail()) { + ABORT_ERROR("BinnedAccumulatorOp: Required parameter 'NumBins' not " + "found in configuration"); + } + + if (NumBins <= 0) { + ABORT_ERROR("BinnedAccumulatorOp: NumBins must be positive, got {}", + NumBins); + } + + // Retrieve input value Field to get dimensions + auto ValueField = Field::get(InputNames[0]); + auto ValueData = ValueField->template getDataArray(); + + // Validate input rank (1D or 2D) + const I4 InputRank = ArrayT::rank; + if (InputRank != 1 && InputRank != 2) { + ABORT_ERROR("BinnedAccumulatorOp: Input must be 1D or 2D, got rank {}", + InputRank); + } + + // Get dimensions + const I4 Dim0 = ValueData.extent(0); + I4 Dim1 = 1; + if (InputRank == 2) { + Dim1 = ValueData.extent(1); + } + + // Store dimensions + HorzSize = Dim0; + VertSize = Dim1; + + // Get dimension info from value field + auto NDims = ValueField->getNumDims(); + std::vector ValueDimNames; + ValueField->getDimNames(ValueDimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = InputNames[0] + "_BinnedAccum"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get input metadata + std::string ValueDescr, ValueUnits, ValueStdName; + ScalarT ValueValidMin, ValueValidMax; + ValueField->getMetadata("Description", ValueDescr); + ValueField->getMetadata("Units", ValueUnits); + ValueField->getMetadata("StdName", ValueStdName); + ValueField->getMetadata("ValidMin", ValueValidMin); + ValueField->getMetadata("ValidMax", ValueValidMax); + + // Create output Field dimensions + std::string NumBinsDimName = "NumBins" + InputNames[1]; + auto NumBinsDim = Dimension::create(NumBinsDimName, NumBins); + + // Replace horizontal dimension with bin dimension + std::vector OutputDimNames; + if (InputRank == 1) { + OutputDimNames = {NumBinsDimName}; + } else { + OutputDimNames = {NumBinsDimName, ValueDimNames[1]}; + } + + // Create output Field + auto OutputField = Field::create( + OutputNames[0], + "Binned accumulation of " + ValueDescr, // Description + ValueUnits, // Units + ValueStdName, // Standard name + ValueValidMin, // Min valid + ValueValidMax, // Max valid + InputRank, // Rank + OutputDimNames // Dimension names + ); + + // Allocate output and local accumulation arrays + if (InputRank == 1) { + OutputData = OutputArrayT("BinnedAccum_out", NumBins); + LocalAccum = OutputArrayT("BinnedAccum_local", NumBins); + } else { + OutputData = OutputArrayT("BinnedAccum_out", NumBins, Dim1); + LocalAccum = OutputArrayT("BinnedAccum_local", NumBins, Dim1); + } + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + } // end constructor + + /// Initializes the operator after Fields are created. Stores mesh pointer, + /// determines the number of owned entities to avoid double-counting halo + /// cells in MPI reductions, and stores MinLayer/MaxLayer arrays from + /// VertCoord for bounding the inner vertical loop. + void initialize(const MachEnv *Env, const HorzMesh *InMesh, + const VertCoord *InVCoord, Config Options) override { + + // Call base class initialization to store Mesh, VCoord, and Comm + AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); + + // Determine index space, set number of owned entities, and MinLayer/MaxLayer + auto ValueField = Field::get(InputNames[0]); + std::vector DimNames; + ValueField->getDimNames(DimNames); + std::string IndexSpaceName = DimNames[0]; + + if (IndexSpaceName == "NCells") { + NOwned = Mesh->NCellsOwned; + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + } else if (IndexSpaceName == "NEdges") { + NOwned = Mesh->NEdgesOwned; + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + } else if (IndexSpaceName == "NVertices") { + NOwned = Mesh->NVerticesOwned; + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + } else { + ABORT_ERROR("BinnedAccumulatorOp: Unknown index space {}", IndexSpaceName); + } + + } // end initialize + + /// Computes the binned accumulation by: + /// 1. Zeroing local accumulation array + /// 2. Accumulating values into bins for owned entities only (excludes halo), + /// applying the Field-attached regional mask if present + /// 3. Performing MPI_Allreduce to sum local accumulations across all ranks + /// 4. Storing global totals in the output array + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Retrieve input fields + auto ValueField = Field::get(InputNames[0]); + auto ValueData = ValueField->template getDataArray(); + auto BinIndexField = Field::get(InputNames[1]); + auto BinIndexData = BinIndexField->template getDataArray(); + + // Apply Field-attached regional mask if present + Array1DI4 MaskData; + bool LocalUseMask = false; + + if (ValueField->hasRegionalMask()) { + MaskData = ValueField->getRegionalMask(); + LocalUseMask = true; + } + + // Dispatch to rank-specific implementation using constexpr to avoid + // instantiating incorrect branches at compile time + if constexpr (ArrayT::rank == 1) { + computeAccum1D(ValueData, BinIndexData, MaskData, LocalUseMask); + } else if constexpr (ArrayT::rank == 2) { + computeAccum2D(ValueData, BinIndexData, MaskData, LocalUseMask); + } else { + ABORT_ERROR("BinnedAccumulatorOp: Unsupported rank {}. " + "Supports 1D and 2D arrays only.", + static_cast(ArrayT::rank)); + } + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Computes binned accumulation for 1D input. Zeroes the local array, + /// then atomically accumulates owned-entity values into bins, and + /// performs MPI_Allreduce to obtain global totals. + void computeAccum1D(const ArrayT &ValueData, + const Array1DI4 &BinIndexData, + const Array1DI4 &MaskData, + bool LocalUseMask) { + + auto LocalAccumData = LocalAccum; + auto LocalNumBins = NumBins; + auto LocalNOwned = NOwned; + + // Zero local accumulation array + parallelFor( + {LocalNumBins}, KOKKOS_LAMBDA(int IBin) { + LocalAccumData(IBin) = 0.0; + }); + + // Accumulate values into bins (only owned entities, not halo) + // Atomic adds handle multiple threads writing to the same bin + parallelFor( + {LocalNOwned}, KOKKOS_LAMBDA(int I) { + I4 BinIdx = BinIndexData(I); + + if (BinIdx >= 0 && BinIdx < LocalNumBins) { + Real Value = static_cast(ValueData(I)); + + if (LocalUseMask) { + Value *= MaskData(I); + } + + Kokkos::atomic_add(&LocalAccumData(BinIdx), Value); + } + }); + + // Copy local accumulation from device to host for MPI reduction + auto LocalAccumHost = createHostMirrorCopy(LocalAccum); + + auto OutputDataHost = createHostMirrorCopy(OutputData); + + // Sum local accumulations across all ranks + MPI_Allreduce(LocalAccumHost.data(), OutputDataHost.data(), NumBins, + MPI_DOUBLE, MPI_SUM, Comm); + + // Copy global totals back to device + deepCopy(OutputData, OutputDataHost); + + } // end computeAccum1D + + /// Computes binned accumulation for 2D input using hierarchical parallelism: + /// outer loop over owned horizontal entities, inner loop over vertical levels. + /// Zeroes the local array, accumulates with atomic adds, then performs + /// MPI_Allreduce to obtain global totals. + void computeAccum2D(const ArrayT &ValueData, + const Array1DI4 &BinIndexData, + const Array1DI4 &MaskData, + bool LocalUseMask) { + + auto LocalAccumData = LocalAccum; + auto LocalNumBins = NumBins; + auto LocalNOwned = NOwned; + auto LocalVertSize = VertSize; + + // Zero local accumulation array + parallelFor( + {LocalNumBins, LocalVertSize}, KOKKOS_LAMBDA(int IBin, int K) { + LocalAccumData(IBin, K) = 0.0; + }); + + // Accumulate values into bins (only owned entities, not halo) + // Hierarchical parallelism: outer over horizontal, inner over vertical + // Inner loop bounded by MinLayer/MaxLayer for partial columns + // Atomic adds handle multiple threads writing to the same bin + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + parallelForOuter( + "BinnedAccum2D", LaunchConfig({LocalNOwned}), + KOKKOS_LAMBDA(int I, const TeamMember &Team) { + I4 BinIdx = BinIndexData(I); + + if (BinIdx >= 0 && BinIdx < LocalNumBins) { + const Real MaskVal = + LocalUseMask ? static_cast(MaskData(I)) : 1.0; + const I4 KMin = LocMinLayer(I); + const I4 KMax = LocMaxLayer(I); + const I4 KRange = vertRange(KMin, KMax); + parallelForInner( + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + Real Value = + static_cast(ValueData(I, K)) * MaskVal; + Kokkos::atomic_add(&LocalAccumData(BinIdx, K), Value); + }); + } + }); + + // Copy local accumulation from device to host for MPI reduction + auto LocalAccumHost = createHostMirrorCopy(LocalAccum); + + auto OutputDataHost = createHostMirrorCopy(OutputData); + + + // Sum local accumulations across all ranks + I4 TotalSize = NumBins * VertSize; + MPI_Allreduce(LocalAccumHost.data(), OutputDataHost.data(), TotalSize, + MPI_DOUBLE, MPI_SUM, Comm); + + // Copy global totals back to device + deepCopy(OutputData, OutputDataHost); + + } // end computeAccum2D + + /// Output data array holding global binned accumulation + OutputArrayT OutputData; + + /// Local accumulation array (before MPI reduction) + OutputArrayT LocalAccum; + + /// Number of spatial bins + I4 NumBins; + + /// Horizontal dimension size (NCells, NEdges, or NVertices) + I4 HorzSize; + + /// Number of owned entities (excludes halo cells for MPI reduction) + I4 NOwned; + + /// Vertical dimension size (NVertLevels or 1 for 1D input) + I4 VertSize; + + /// Min active layer index for each horizontal point (for 2D input) + Array1DI4 MinLayer; + + /// Max active layer index for each horizontal point (for 2D input) + Array1DI4 MaxLayer; + +}; // end class BinnedAccumulatorOp + +} // end namespace OMEGA + +#endif From 1aae71713a785035b3f3b906eac62daae9fb4758 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sun, 26 Jul 2026 20:20:17 -0400 Subject: [PATCH 06/35] Add PrefixSumOp --- .../src/analysis/operators/BinaryMultiplyOp.h | 76 +++- .../analysis/operators/BinnedAccumulatorOp.h | 66 ++- .../src/analysis/operators/PrefixSumOp.h | 404 ++++++++++++++++++ 3 files changed, 494 insertions(+), 52 deletions(-) create mode 100644 components/omega/src/analysis/operators/PrefixSumOp.h diff --git a/components/omega/src/analysis/operators/BinaryMultiplyOp.h b/components/omega/src/analysis/operators/BinaryMultiplyOp.h index 32c5677ff016..26a679479f97 100644 --- a/components/omega/src/analysis/operators/BinaryMultiplyOp.h +++ b/components/omega/src/analysis/operators/BinaryMultiplyOp.h @@ -50,7 +50,7 @@ template class BinaryMultiplyOp : public AnalysisOperator { /// dimensions and metadata, validates that they match, creates output /// Field for the product, allocates output data array, and registers /// the output Field. The output Field name is constructed as - /// Field1Name + "_BinaryMultiply(" + Field2Name + ")_Product". + /// Field1Name + "_" + Field2Name + "_Product". BinaryMultiplyOp(const std::vector &UpstreamNames, ///< [in] input field names Config Options ///< [in] operator config @@ -184,13 +184,41 @@ template class BinaryMultiplyOp : public AnalysisOperator { /// Initializes the operator after all Fields exist. Determines the index /// space (cells, edges, or vertices) from the first input field's horizontal /// dimension name and stores the appropriate MinLayer/MaxLayer arrays - /// from VertCoord for bounding the inner vertical loop. + /// from VertCoord for bounding the inner vertical loop. Also determines + /// whether the horizontal dimension is mesh-distributed and stores the + /// appropriate owned count for MPI-correct parallel loops. void initialize(const MachEnv *Env, const HorzMesh *InMesh, const VertCoord *InVCoord, Config Options) override { AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); constexpr I4 InputRank = ArrayT::rank; + + // For 1D arrays, check if it's a mesh dimension to determine owned count + if constexpr (InputRank == 1) { + auto Field1 = Field::get(InputNames[0]); + std::vector DimNames; + Field1->getDimNames(DimNames); + std::string IndexSpaceName = DimNames[0]; + + IsMeshDimension = + (IndexSpaceName == "NCells" || IndexSpaceName == "NEdges" || + IndexSpaceName == "NVertices"); + + if (IsMeshDimension) { + if (IndexSpaceName == "NCells") { + NHorizOwned = Mesh->NCellsOwned; + } else if (IndexSpaceName == "NEdges") { + NHorizOwned = Mesh->NEdgesOwned; + } else if (IndexSpaceName == "NVertices") { + NHorizOwned = Mesh->NVerticesOwned; + } + } else { + ABORT_ERROR("BinaryMultiplyOp: Unknown index space {}", + IndexSpaceName); + } + } + if constexpr (InputRank > 1) { auto Field1 = Field::get(InputNames[0]); std::vector DimNames; @@ -198,15 +226,22 @@ template class BinaryMultiplyOp : public AnalysisOperator { // Horizontal dimension is 2nd-to-last for 2D/3D std::string IndexSpaceName = DimNames[InputRank - 2]; + IsMeshDimension = + (IndexSpaceName == "NCells" || IndexSpaceName == "NEdges" || + IndexSpaceName == "NVertices"); + if (IndexSpaceName == "NCells") { - MinLayer = VCoord->MinLayerCell; - MaxLayer = VCoord->MaxLayerCell; + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + NHorizOwned = Mesh->NCellsOwned; } else if (IndexSpaceName == "NEdges") { - MinLayer = VCoord->MinLayerEdgeBot; - MaxLayer = VCoord->MaxLayerEdgeTop; + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + NHorizOwned = Mesh->NEdgesOwned; } else if (IndexSpaceName == "NVertices") { - MinLayer = VCoord->MinLayerVertexBot; - MaxLayer = VCoord->MaxLayerVertexTop; + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + NHorizOwned = Mesh->NVerticesOwned; } else { ABORT_ERROR("BinaryMultiplyOp: Unknown index space {}", IndexSpaceName); @@ -240,7 +275,7 @@ template class BinaryMultiplyOp : public AnalysisOperator { // 1D case: horizontal field only, no vertical structure auto Data2 = Field2->template getDataArray(); parallelFor( - {NHorizDim}, KOKKOS_LAMBDA(int IHoriz) { + {NHorizOwned}, KOKKOS_LAMBDA(int IHoriz) { Output(IHoriz) = static_cast(static_cast(Data1(IHoriz)) * static_cast(Data2(IHoriz))); @@ -257,7 +292,7 @@ template class BinaryMultiplyOp : public AnalysisOperator { // Vertical expansion: Data2 is 1D, replicate across vertical layers auto Data2_1D = Field2->template getDataArray(); parallelForOuter( - "BinaryMultiply2D_VertExpand", LaunchConfig({NHorizDim}), + "BinaryMultiply2D_VertExpand", LaunchConfig({NHorizOwned}), KOKKOS_LAMBDA(int IHoriz, const TeamMember &Team) { const Real Data2Val = static_cast(Data2_1D(IHoriz)); const I4 KMin = LocMinLayer(IHoriz); @@ -265,8 +300,8 @@ template class BinaryMultiplyOp : public AnalysisOperator { const I4 KRange = vertRange(KMin, KMax); parallelForInner( Team, KRange, INNER_LAMBDA(int KIdx) { - const I4 K = KMin + KIdx; - Output(IHoriz, K) = static_cast( + const I4 K = KMin + KIdx; + Output(IHoriz, K) = static_cast( static_cast(Data1(IHoriz, K)) * Data2Val); }); }); @@ -274,7 +309,7 @@ template class BinaryMultiplyOp : public AnalysisOperator { // Same rank: element-wise multiplication auto Data2 = Field2->template getDataArray(); parallelForOuter( - "BinaryMultiply2D", LaunchConfig({NHorizDim}), + "BinaryMultiply2D", LaunchConfig({NHorizOwned}), KOKKOS_LAMBDA(int IHoriz, const TeamMember &Team) { const I4 KMin = LocMinLayer(IHoriz); const I4 KMax = LocMaxLayer(IHoriz); @@ -301,7 +336,8 @@ template class BinaryMultiplyOp : public AnalysisOperator { // Vertical expansion: Data2 is 1D, replicate across vertical layers auto Data2_1D = Field2->template getDataArray(); parallelForOuter( - "BinaryMultiply3D_VertExpand", LaunchConfig({Dim0, NHorizDim}), + "BinaryMultiply3D_VertExpand", + LaunchConfig({Dim0, NHorizOwned}), KOKKOS_LAMBDA(int I0, int IHoriz, const TeamMember &Team) { const Real Data2Val = static_cast(Data2_1D(IHoriz)); const I4 KMin = LocMinLayer(IHoriz); @@ -309,8 +345,8 @@ template class BinaryMultiplyOp : public AnalysisOperator { const I4 KRange = vertRange(KMin, KMax); parallelForInner( Team, KRange, INNER_LAMBDA(int KIdx) { - const I4 K = KMin + KIdx; - Output(I0, IHoriz, K) = static_cast( + const I4 K = KMin + KIdx; + Output(I0, IHoriz, K) = static_cast( static_cast(Data1(I0, IHoriz, K)) * Data2Val); }); @@ -319,7 +355,7 @@ template class BinaryMultiplyOp : public AnalysisOperator { // Same rank: element-wise multiplication auto Data2 = Field2->template getDataArray(); parallelForOuter( - "BinaryMultiply3D", LaunchConfig({Dim0, NHorizDim}), + "BinaryMultiply3D", LaunchConfig({Dim0, NHorizOwned}), KOKKOS_LAMBDA(int I0, int IHoriz, const TeamMember &Team) { const I4 KMin = LocMinLayer(IHoriz); const I4 KMax = LocMaxLayer(IHoriz); @@ -363,9 +399,15 @@ template class BinaryMultiplyOp : public AnalysisOperator { /// Number of points in horizontal dimension (cells, edges, or vertices) I4 NHorizDim; + /// Number of owned points in horizontal dimension (for MPI correctness) + I4 NHorizOwned; + /// Vertical size of the array I4 NVertSize; + /// Whether the horizontal dimension is a mesh-distributed dimension + bool IsMeshDimension; + /// Whether Field2 is 1D and should be replicated across vertical dimension bool IsVerticalExpansion; diff --git a/components/omega/src/analysis/operators/BinnedAccumulatorOp.h b/components/omega/src/analysis/operators/BinnedAccumulatorOp.h index 8e11460118b9..728b1232cba5 100644 --- a/components/omega/src/analysis/operators/BinnedAccumulatorOp.h +++ b/components/omega/src/analysis/operators/BinnedAccumulatorOp.h @@ -1,7 +1,7 @@ #ifndef OMEGA_BINNEDACCUMULATOROP_H #define OMEGA_BINNEDACCUMULATOROP_H -//===-- analysis/operators/BinnedAccumulatorOp.h - BinnedAccumulatorOp -*- C++ -*-===// +//===-- analysis/operators/BinnedAccumulatorOp.h ----------------*- C++ -*-===// // /// \file /// \brief Defines the BinnedAccumulatorOp operator for binned accumulation @@ -51,9 +51,10 @@ template class BinnedAccumulatorOp : public AnalysisOperator { using ScalarT = typename ArrayT::non_const_value_type; /// Output array type depends on input rank: - /// For 2D input (e.g., NCells x NVertLevels), output is 2D (NumBins x NVertLevels) - /// For 1D input (e.g., NCells), output is 1D (NumBins) - using OutputArrayT = std::conditional_t; + /// For 2D input (e.g., NCells x NVertLevels), output is 2D (NumBins x + /// NVertLevels) For 1D input (e.g., NCells), output is 1D (NumBins) + using OutputArrayT = + std::conditional_t; /// Constructs a BinnedAccumulatorOp operator. Reads the required NumBins /// configuration parameter, creates the output Field for binned @@ -131,7 +132,7 @@ template class BinnedAccumulatorOp : public AnalysisOperator { // Create output Field dimensions std::string NumBinsDimName = "NumBins" + InputNames[1]; - auto NumBinsDim = Dimension::create(NumBinsDimName, NumBins); + auto NumBinsDim = Dimension::create(NumBinsDimName, NumBins); // Replace horizontal dimension with bin dimension std::vector OutputDimNames; @@ -142,16 +143,16 @@ template class BinnedAccumulatorOp : public AnalysisOperator { } // Create output Field - auto OutputField = Field::create( - OutputNames[0], - "Binned accumulation of " + ValueDescr, // Description - ValueUnits, // Units - ValueStdName, // Standard name - ValueValidMin, // Min valid - ValueValidMax, // Max valid - InputRank, // Rank - OutputDimNames // Dimension names - ); + auto OutputField = + Field::create(OutputNames[0], + "Binned accumulation of " + ValueDescr, // Description + ValueUnits, // Units + ValueStdName, // Standard name + ValueValidMin, // Min valid + ValueValidMax, // Max valid + InputRank, // Rank + OutputDimNames // Dimension names + ); // Allocate output and local accumulation arrays if (InputRank == 1) { @@ -177,7 +178,8 @@ template class BinnedAccumulatorOp : public AnalysisOperator { // Call base class initialization to store Mesh, VCoord, and Comm AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); - // Determine index space, set number of owned entities, and MinLayer/MaxLayer + // Determine index space, set number of owned entities, and + // MinLayer/MaxLayer auto ValueField = Field::get(InputNames[0]); std::vector DimNames; ValueField->getDimNames(DimNames); @@ -196,7 +198,8 @@ template class BinnedAccumulatorOp : public AnalysisOperator { MinLayer = VCoord->MinLayerVertexBot; MaxLayer = VCoord->MaxLayerVertexTop; } else { - ABORT_ERROR("BinnedAccumulatorOp: Unknown index space {}", IndexSpaceName); + ABORT_ERROR("BinnedAccumulatorOp: Unknown index space {}", + IndexSpaceName); } } // end initialize @@ -247,10 +250,8 @@ template class BinnedAccumulatorOp : public AnalysisOperator { /// Computes binned accumulation for 1D input. Zeroes the local array, /// then atomically accumulates owned-entity values into bins, and /// performs MPI_Allreduce to obtain global totals. - void computeAccum1D(const ArrayT &ValueData, - const Array1DI4 &BinIndexData, - const Array1DI4 &MaskData, - bool LocalUseMask) { + void computeAccum1D(const ArrayT &ValueData, const Array1DI4 &BinIndexData, + const Array1DI4 &MaskData, bool LocalUseMask) { auto LocalAccumData = LocalAccum; auto LocalNumBins = NumBins; @@ -258,9 +259,8 @@ template class BinnedAccumulatorOp : public AnalysisOperator { // Zero local accumulation array parallelFor( - {LocalNumBins}, KOKKOS_LAMBDA(int IBin) { - LocalAccumData(IBin) = 0.0; - }); + {LocalNumBins}, + KOKKOS_LAMBDA(int IBin) { LocalAccumData(IBin) = 0.0; }); // Accumulate values into bins (only owned entities, not halo) // Atomic adds handle multiple threads writing to the same bin @@ -294,13 +294,11 @@ template class BinnedAccumulatorOp : public AnalysisOperator { } // end computeAccum1D /// Computes binned accumulation for 2D input using hierarchical parallelism: - /// outer loop over owned horizontal entities, inner loop over vertical levels. - /// Zeroes the local array, accumulates with atomic adds, then performs - /// MPI_Allreduce to obtain global totals. - void computeAccum2D(const ArrayT &ValueData, - const Array1DI4 &BinIndexData, - const Array1DI4 &MaskData, - bool LocalUseMask) { + /// outer loop over owned horizontal entities, inner loop over vertical + /// levels. Zeroes the local array, accumulates with atomic adds, then + /// performs MPI_Allreduce to obtain global totals. + void computeAccum2D(const ArrayT &ValueData, const Array1DI4 &BinIndexData, + const Array1DI4 &MaskData, bool LocalUseMask) { auto LocalAccumData = LocalAccum; auto LocalNumBins = NumBins; @@ -309,9 +307,8 @@ template class BinnedAccumulatorOp : public AnalysisOperator { // Zero local accumulation array parallelFor( - {LocalNumBins, LocalVertSize}, KOKKOS_LAMBDA(int IBin, int K) { - LocalAccumData(IBin, K) = 0.0; - }); + {LocalNumBins, LocalVertSize}, + KOKKOS_LAMBDA(int IBin, int K) { LocalAccumData(IBin, K) = 0.0; }); // Accumulate values into bins (only owned entities, not halo) // Hierarchical parallelism: outer over horizontal, inner over vertical @@ -345,7 +342,6 @@ template class BinnedAccumulatorOp : public AnalysisOperator { auto OutputDataHost = createHostMirrorCopy(OutputData); - // Sum local accumulations across all ranks I4 TotalSize = NumBins * VertSize; MPI_Allreduce(LocalAccumHost.data(), OutputDataHost.data(), TotalSize, diff --git a/components/omega/src/analysis/operators/PrefixSumOp.h b/components/omega/src/analysis/operators/PrefixSumOp.h new file mode 100644 index 000000000000..624f9940071f --- /dev/null +++ b/components/omega/src/analysis/operators/PrefixSumOp.h @@ -0,0 +1,404 @@ +#ifndef OMEGA_PREFIXSUMOP_H +#define OMEGA_PREFIXSUMOP_H + +//===-- analysis/operators/PrefixSumOp.h - PrefixSumOp ----------*- C++ -*-===// +// +/// \file +/// \brief Defines the PrefixSumOp operator for cumulative summation +/// +/// PrefixSumOp computes the cumulative sum (prefix sum, or scan) of an input +/// field along a specified dimension. The operator supports both forward +/// (inclusive scan from start to end) and reverse (inclusive scan from end to +/// start) directions. +/// +/// Mathematical formulation: +/// - Forward scan: Output[i] = sum(Input[0] + Input[1] + ... + Input[i]) +/// - Reverse scan: Output[i] = sum(Input[i] + Input[i+1] + ... + Input[n-1]) +/// +/// This is particularly useful for vertical integration in ocean models, where +/// the MOC (Meridional Overturning Circulation) stream function requires +/// cumulative integration of transport from bottom to top. The reverse scan +/// direction is typically used for this case, starting from the ocean bottom +/// (maximum depth index) and integrating upward to the surface. +/// +/// The operator is templated on the Kokkos array type (ArrayT) of the input +/// field. Currently supports 2D arrays (e.g., latitude bins × depth levels). +/// The output has the same shape and dimensions as the input. +/// +/// Configuration: +/// - Dimension: The dimension along which to compute cumulative sum (required) +/// 0 = first dimension, 1 = second dimension +/// Note: For 2D arrays with mesh dimensions (NCells, NEdges, +/// NVertices), scanning along dimension 0 is not allowed in +/// distributed (MPI) configurations. Horizontal scans are only +/// supported for non-distributed dimensions (e.g., bins). +/// - Reverse: If true, scan from end to start; if false, scan from start to +/// end (default: false) +/// +/// Example usage in operator chain: +/// \code +/// BinnedTransport_PrefixSum +/// \endcode +/// where PrefixSum performs cumulative vertical integration for MOC, typically +/// with Dimension=1 and Reverse=true to integrate from ocean bottom to top. +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" + +namespace OMEGA { + +/// PrefixSumOp computes the cumulative sum of an input field along a specified +/// dimension. Supports forward and reverse directions for flexibility. The +/// output array has the same shape as the input. Currently optimized for 2D +/// arrays common in MOC calculations. +template class PrefixSumOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type - same as input array type + using OutputArrayT = ArrayT; + + /// Constructs a PrefixSumOp operator. Creates output Field with the same + /// dimensions as the input field, allocates output data array, and + /// registers the output Field in the Field registry. The output Field name + /// is constructed as InputName + "_PrefixSum". Reads the dimension and + /// direction parameters from the Options config. + PrefixSumOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("PrefixSum") { + + // Store input field names + InputNames = UpstreamNames; + + // Read required dimension parameter from configuration + Error Err = Options.get("Dimension", ScanDimension); + if (Err.isFail()) { + ABORT_ERROR("PrefixSumOp: Required parameter 'Dimension' not found in " + "configuration"); + } + + // Read optional reverse parameter (default: false = forward scan) + Err = Options.get("Reverse", ReverseDirection); + if (Err.isFail()) { + ReverseDirection = false; // Default to forward scan + } + + // Retrieve input Field to get dimensions and metadata + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Validate dimension parameter + if (ScanDimension < 0 || ScanDimension >= static_cast(ArrayT::rank)) { + ABORT_ERROR("PrefixSumOp: Dimension {} out of range for array rank {}", + ScanDimension, static_cast(ArrayT::rank)); + } + + // Get dimension info + auto NDims = InputField->getNumDims(); + std::vector DimNames; + InputField->getDimNames(DimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = InputNames[0] + "_PrefixSum"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get input metadata + std::string InputDescr, InputUnits, InputStdName; + ScalarT InputValidMin, InputValidMax; + InputField->getMetadata("Description", InputDescr); + InputField->getMetadata("Units", InputUnits); + InputField->getMetadata("StdName", InputStdName); + InputField->getMetadata("ValidMin", InputValidMin); + InputField->getMetadata("ValidMax", InputValidMax); + + // Create output Field with same dimensions as input + auto OutputField = + Field::create(OutputNames[0], + "Cumulative sum of " + InputDescr, // Description + InputUnits, // Units + InputStdName, // Standard name + InputValidMin, // Min valid + InputValidMax, // Max valid + NDims, // Rank + DimNames // Dimension names + ); + + // Allocate output data array matching input layout + OutputData = OutputArrayT(OutputNames[0] + "_out", InputData.layout()); + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + } // end constructor + + /// Initializes the operator after all Fields exist. For 2D arrays, + /// examines the first dimension name to determine if it's a mesh dimension + /// (NCells, NEdges, NVertices) or a non-mesh dimension (e.g., bins). For + /// mesh dimensions, stores the appropriate MinLayer/MaxLayer arrays from + /// VertCoord for bounding the inner vertical loop when scanning along the + /// vertical dimension (dimension 1). For non-mesh dimensions (bins), no + /// MinLayer/MaxLayer bounds are needed. The dimension name is used in + /// compute() to determine whether to use owned count (mesh dimensions) or + /// full extent (non-mesh dimensions) for MPI-correct parallel loops. + void initialize(const MachEnv *Env, const HorzMesh *InMesh, + const VertCoord *InVCoord, Config Options) override { + + AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); + + // Only need MinLayer/MaxLayer and dimension info for 2D arrays + if constexpr (ArrayT::rank == 2) { + auto Field1 = Field::get(InputNames[0]); + std::vector DimNames; + Field1->getDimNames(DimNames); + // First dimension can be mesh (NCells/NEdges/NVertices) or non-mesh + // (bins) + Dim0Name = DimNames[0]; + + if (Dim0Name == "NCells") { + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + IsMeshDimension = true; + } else if (Dim0Name == "NEdges") { + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + IsMeshDimension = true; + } else if (Dim0Name == "NVertices") { + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + IsMeshDimension = true; + } else { + // For non-mesh dimensions (e.g., bins), MinLayer/MaxLayer remain + // uninitialized + IsMeshDimension = false; + } + } + + } // end initialize + + /// Computes the cumulative sum along the specified dimension. For 2D + /// arrays, uses hierarchical parallelism with parallelForOuter and + /// parallelScanInner. Supports both forward (start to end) and reverse + /// (end to start) directions. Updates output data, timestamp, and computed + /// flag. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Retrieve input Field and extract data array + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Compute cumulative sum based on array rank using constexpr to prevent + // instantiation of incorrect branches at compile time + if constexpr (ArrayT::rank == 2) { + compute2D(InputData); + } else { + ABORT_ERROR( + "PrefixSumOp: Currently only 2D arrays are supported. Rank: {}", + static_cast(ArrayT::rank)); + } + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Computes cumulative sum for 2D arrays using hierarchical parallelism. + /// Uses parallelForOuter for the non-scan dimension and parallelScanInner + /// for the scan dimension. For vertical scans (dimension 1), bounds the + /// loop with MinLayer/MaxLayer to handle partial columns. Horizontal scans + /// (dimension 0) are only allowed for non-distributed dimensions like bins. + void compute2D(const ArrayT &InputData) { + + // Check if scanning along distributed dimension (not allowed) + if (ScanDimension == 0 && (Dim0Name == "NCells" || Dim0Name == "NEdges" || + Dim0Name == "NVertices")) { + ABORT_ERROR("PrefixSumOp: Cannot scan along distributed dimension {}. " + "Horizontal scans only supported for non-distributed " + "dimensions (e.g., bins).", + Dim0Name); + } + + // Determine loop extent for dimension 0 + // For mesh dimensions, use owned count; for non-mesh dimensions, use full + // extent + I4 N0; + if (IsMeshDimension) { + if (Dim0Name == "NCells") { + N0 = Mesh->NCellsOwned; + } else if (Dim0Name == "NEdges") { + N0 = Mesh->NEdgesOwned; + } else if (Dim0Name == "NVertices") { + N0 = Mesh->NVerticesOwned; + } else { + N0 = InputData.extent(0); // fallback + } + } else { + N0 = InputData.extent( + 0); // bins or other non-distributed dims use full extent + } + const I4 N1 = InputData.extent(1); + + auto LocalOutput = OutputData; + auto LocalDim = ScanDimension; + auto LocalReverse = ReverseDirection; + + if (LocalDim == 0) { + // Scan along first dimension (horizontal bins) + // Outer loop over second dimension (vertical), inner scan over first + if (LocalReverse) { + // Reverse: scan from end to start + parallelForOuter( + "PrefixSum_Dim0_Reverse", LaunchConfig({N1}), + KOKKOS_LAMBDA(int j, const TeamMember &Team) { + parallelScanInner( + Team, N0, + INNER_LAMBDA(int iRev, Real &Accum, bool IsFinal) { + const int i = N0 - 1 - iRev; + Accum += static_cast(InputData(i, j)); + if (IsFinal) { + LocalOutput(i, j) = static_cast(Accum); + } + }); + }); + } else { + // Forward: scan from start to end + parallelForOuter( + "PrefixSum_Dim0_Forward", LaunchConfig({N1}), + KOKKOS_LAMBDA(int j, const TeamMember &Team) { + parallelScanInner( + Team, N0, + INNER_LAMBDA(int i, Real &Accum, bool IsFinal) { + Accum += static_cast(InputData(i, j)); + if (IsFinal) { + LocalOutput(i, j) = static_cast(Accum); + } + }); + }); + } + } else { // LocalDim == 1 + // Scan along second dimension (vertical) + // Outer loop over first dimension (bins/etc), inner scan over vertical + + if (IsMeshDimension) { + // For mesh dimensions, use MinLayer/MaxLayer to bound vertical + // loops for partial columns + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + + if (LocalReverse) { + // Reverse: scan from end to start + // Iterate from KMax down to KMin, accumulating as we go + // Result at K = sum from K to KMax + parallelForOuter( + "PrefixSum_Dim1_Reverse", LaunchConfig({N0}), + KOKKOS_LAMBDA(int i, const TeamMember &Team) { + const I4 KMin = LocMinLayer(i); + const I4 KMax = LocMaxLayer(i); + const I4 KRange = vertRange(KMin, KMax); + parallelScanInner( + Team, KRange, + INNER_LAMBDA(int KIdx, Real &Accum, bool IsFinal) { + // Iterate from KMax down: KIdx 0→KRange-1 maps to + // K KMax→KMin + const I4 K = KMax - KIdx; + Accum += static_cast(InputData(i, K)); + if (IsFinal) { + LocalOutput(i, K) = static_cast(Accum); + } + }); + }); + } else { + // Forward: scan from start to end (top to bottom for vertical) + parallelForOuter( + "PrefixSum_Dim1_Forward", LaunchConfig({N0}), + KOKKOS_LAMBDA(int i, const TeamMember &Team) { + const I4 KMin = LocMinLayer(i); + const I4 KMax = LocMaxLayer(i); + const I4 KRange = vertRange(KMin, KMax); + parallelScanInner( + Team, KRange, + INNER_LAMBDA(int KIdx, Real &Accum, bool IsFinal) { + // Map KIdx to actual layer: 0→KMin, KRange-1→KMax + const I4 K = KMin + KIdx; + Accum += static_cast(InputData(i, K)); + if (IsFinal) { + LocalOutput(i, K) = static_cast(Accum); + } + }); + }); + } + } else { + // For non-mesh dimensions (e.g., bins), scan over full extent + if (LocalReverse) { + // Reverse: scan from end to start + parallelForOuter( + "PrefixSum_Dim1_Reverse_Full", LaunchConfig({N0}), + KOKKOS_LAMBDA(int i, const TeamMember &Team) { + parallelScanInner( + Team, N1, + INNER_LAMBDA(int KIdx, Real &Accum, bool IsFinal) { + // Iterate from N1-1 down to 0 + const I4 K = N1 - 1 - KIdx; + Accum += static_cast(InputData(i, K)); + if (IsFinal) { + LocalOutput(i, K) = static_cast(Accum); + } + }); + }); + } else { + // Forward: scan from start to end + parallelForOuter( + "PrefixSum_Dim1_Forward_Full", LaunchConfig({N0}), + KOKKOS_LAMBDA(int i, const TeamMember &Team) { + parallelScanInner( + Team, N1, + INNER_LAMBDA(int KIdx, Real &Accum, bool IsFinal) { + // KIdx maps directly to K + const I4 K = KIdx; + Accum += static_cast(InputData(i, K)); + if (IsFinal) { + LocalOutput(i, K) = static_cast(Accum); + } + }); + }); + } + } + } + + } // end compute2D + + /// Output data array holding the cumulative sum + OutputArrayT OutputData; + + /// The dimension along which to perform cumulative sum (0 or 1 for 2D) + I4 ScanDimension; + + /// Whether to scan in reverse direction (true) or forward (false) + bool ReverseDirection; + + /// Min active layer index for each horizontal point (only for 2D arrays) + Array1DI4 MinLayer; + + /// Max active layer index for each horizontal point (only for 2D arrays) + Array1DI4 MaxLayer; + + /// Name of the first dimension (only for 2D arrays) to determine owned count + std::string Dim0Name; + + /// Flag indicating whether the first dimension is a mesh dimension (NCells, + /// NEdges, NVertices) If true, MinLayer/MaxLayer are used; if false, full + /// extent is used + bool IsMeshDimension = false; + +}; // end class PrefixSumOp + +} // end namespace OMEGA + +#endif From 54e14f10983e748503b9fa8c7b43f8de5cc7dfd0 Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sun, 26 Jul 2026 21:08:34 -0400 Subject: [PATCH 07/35] Add ScalarMultiplyOp --- .../src/analysis/operators/ScalarMultiplyOp.h | 184 ++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 components/omega/src/analysis/operators/ScalarMultiplyOp.h diff --git a/components/omega/src/analysis/operators/ScalarMultiplyOp.h b/components/omega/src/analysis/operators/ScalarMultiplyOp.h new file mode 100644 index 000000000000..1a6aed14509d --- /dev/null +++ b/components/omega/src/analysis/operators/ScalarMultiplyOp.h @@ -0,0 +1,184 @@ +#ifndef OMEGA_SCALARMULTIPLYOP_H +#define OMEGA_SCALARMULTIPLYOP_H + +//===-- analysis/operators/ScalarMultiplyOp.h -------------------*- C++ -*-===// +// +/// \file +/// \brief Defines the ScalarMultiplyOp operator for scalar multiplication +/// +/// ScalarMultiplyOp multiplies all elements of an input field by a +/// user-specified scalar value. The operator is templated on the Kokkos array +/// type (ArrayT) of the input field, supporting 1D, 2D, and 3D+ fields. The +/// output has the same shape and dimensions as the input. +/// +/// This operator is particularly useful for unit conversions (e.g., m³/s to +/// Sverdrups for ocean transport) and for scaling derived quantities. +/// +/// Configuration: +/// - Scalar: The multiplicative factor (specified in parentheses in operator +/// chain) +/// - InPlace: If true, modifies input array directly (default: false) +/// +/// Example usage in operator chain: +/// \code +/// Field_ScalarMultiply(1.0e-6) +/// \endcode +/// where ScalarMultiply(1.0e-6) multiplies the Field by 1.0e-6. +/// This is useful for unit conversions (e.g., m³/s to Sv). +/// When InPlace=true, no additional memory is allocated; the input array is +/// modified directly and the output Field references the same data. +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" + +namespace OMEGA { + +/// ScalarMultiplyOp multiplies all elements of an input field by a scalar +/// value. The operator handles arrays of any rank (1D, 2D, 3D+) and preserves +/// the input dimensions in the output. The scalar multiplication is performed +/// element-wise using Kokkos parallel operations. Output type matches input +/// type unless input is integral, in which case output is Real. +template class ScalarMultiplyOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type - same as input array type + using OutputArrayT = ArrayT; + + /// Constructs a ScalarMultiplyOp operator. Creates output Field with the + /// same dimensions as the input field, allocates output data array, and + /// registers the output Field in the Field registry. The output Field name + /// is constructed as InputName + "_ScalarMultiply". Reads the scalar + /// multiplicative factor from the Options config (which is set by the parser + /// from the parenthesized argument in the operator chain). + ScalarMultiplyOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("ScalarMultiply") { + + // Store input field names + InputNames = UpstreamNames; + + // Read required scalar parameter from configuration + Error Err = Options.get("Scalar", Scalar); + + if (Err.isFail()) { + ABORT_ERROR("ScalarMultiplyOp: Required parameter 'Scalar' not found " + "in configuration"); + } + + // Read optional InPlace parameter (default: false) + Err = Options.get("InPlace", InPlace); + if (Err.isFail()) { + InPlace = false; // Default to allocating new array + } + + // Retrieve input Field to get dimensions and metadata + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Get dimension info + auto NDims = InputField->getNumDims(); + std::vector DimNames; + InputField->getDimNames(DimNames); + + // Construct output field name and set instance name + std::string OutputFieldName = InputNames[0] + "_ScalarMultiply"; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get input metadata + std::string InputDescr, InputUnits, InputStdName; + ScalarT InputValidMin, InputValidMax, InputFillValue; + InputField->getMetadata("Description", InputDescr); + InputField->getMetadata("StdName", InputStdName); + InputField->getMetadata("ValidMin", InputValidMin); + InputField->getMetadata("ValidMax", InputValidMax); + + // Create output Field with same dimensions as input + auto OutputField = + Field::create(OutputNames[0], + InputDescr + " multiplied by " + + std::to_string(Scalar), // Description + "", // Units + InputStdName, // Standard name + InputValidMin, // Min valid + InputValidMax, // Max valid + NDims, // Rank + DimNames // Dimension names + ); + + // Store array size for parallel iteration + ArraySize = static_cast(InputData.size()); + + // Handle InPlace vs new allocation + if (InPlace) { + // InPlace: output Field references the same array as input + // This modifies the input data directly + OutputData = InputData; + } else { + // Allocate new output data array matching input layout + OutputData = OutputArrayT(OutputNames[0] + "_out", InputData.layout()); + } + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + // Propagate regional mask from input to output if present + if (InputField->hasRegionalMask()) { + OutputField->setRegionalMask(InputField->getRegionalMask()); + } + + } // end constructor + + /// Computes the scalar multiplication by retrieving input data and + /// performing element-wise multiplication using Kokkos parallel_for. The + /// operation uses flat indexing to handle arrays of any rank efficiently. + /// Updates output data, timestamp, and computed flag. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Create local scope reference to output array for kernel capture + OMEGA_SCOPE(LocOutputData, OutputData); + + // Retrieve input Field and extract data array + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Perform element-wise multiplication using flat indexing + // This works for arrays of any rank + auto LocalScalar = Scalar; + parallelFor( + {ArraySize}, KOKKOS_LAMBDA(const int FlatIdx) { + LocOutputData.data()[FlatIdx] = + InputData.data()[FlatIdx] * LocalScalar; + }); + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Output data array holding the scaled field values + OutputArrayT OutputData; + + /// The scalar multiplier read from configuration + Real Scalar; + + /// Whether to modify input array in-place (true) or allocate new array + /// (false) + bool InPlace; + + /// Total size of the array for flat indexing + I4 ArraySize; + +}; // end class ScalarMultiplyOp + +} // end namespace OMEGA + +#endif From e748f7a0b391e79cc69a70f286a95701263c1b6e Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Sun, 26 Jul 2026 22:27:24 -0400 Subject: [PATCH 08/35] Add ExtractRegionOp.h --- .../src/analysis/operators/ExtractRegionOp.h | 382 ++++++++++++++++++ 1 file changed, 382 insertions(+) create mode 100644 components/omega/src/analysis/operators/ExtractRegionOp.h diff --git a/components/omega/src/analysis/operators/ExtractRegionOp.h b/components/omega/src/analysis/operators/ExtractRegionOp.h new file mode 100644 index 000000000000..8b24aca9511e --- /dev/null +++ b/components/omega/src/analysis/operators/ExtractRegionOp.h @@ -0,0 +1,382 @@ +#ifndef OMEGA_EXTRACTREGIONOP_H +#define OMEGA_EXTRACTREGIONOP_H + +//===-- analysis/operators/ExtractRegionOp.h - ExtractRegionOp -*- C++ -*-===// +// +/// \file +/// \brief Defines the ExtractRegionOp operator for applying regional masks +/// +/// ExtractRegionOp applies a named regional mask to a field by multiplying +/// the field values by the mask. The mask must already exist in the Field +/// registry as an Array1DI4 field. This operator enables regional statistics +/// by multiplying field values by mask values (where mask == 0 excludes points +/// by leaving output unchanged), which can then be processed by spatial +/// reduction operators that use the attached regional mask. +/// +/// The operator multiplies each horizontal location by its mask value: +/// Output(i, k, ...) = Input(i, k, ...) * Mask(i) +/// +/// where Mask(i) is typically 0 (exclude) or 1 (include). +/// +/// Configuration: +/// - MaskName: Name of the mask field in the Field registry (required) +/// +/// Example usage in operator chain: +/// \code +/// Temperature_ExtractRegion(Atlantic)_SpatialMean +/// \endcode +/// where `ExtractRegion(Atlantic)` applies the "Atlantic" mask to the +/// Temperature field, enabling regional spatial statistics. +/// +/// The mask field should be Array1DI4 with dimension matching the horizontal +/// dimension of the input field (NCells, NEdges, or NVertices). Values are: +/// - 0: exclude this horizontal location (data not copied to output) +/// - 1: include this horizontal location (data multiplied by 1 and copied) +/// +//===----------------------------------------------------------------------===// + +#include "AnalysisOperator.h" + +namespace OMEGA { + +/// ExtractRegionOp applies a named regional mask from the Field registry to +/// a field by multiplying field values with the mask. The operator copies +/// masked data to the output field and attaches the mask to enable downstream +/// spatial reduction operators to exclude masked-out points. +template class ExtractRegionOp : public AnalysisOperator { + public: + /// Scalar type extracted from the input array type + using ScalarT = typename ArrayT::non_const_value_type; + + /// Output array type - same as input array type + using OutputArrayT = ArrayT; + + /// Constructs an ExtractRegionOp operator. Reads the mask field name from + /// config, retrieves the mask from the Field registry, creates an output + /// Field matching the input dimensions, and attaches the regional mask to + /// the output Field. The output Field name is constructed as + /// InputName + "_" + MaskName (without "Mask" suffix if present). + ExtractRegionOp(const std::vector + &UpstreamNames, ///< [in] input field names + Config Options ///< [in] operator config + ) + : AnalysisOperator("ExtractRegion") { + + // Store input field names + InputNames = UpstreamNames; + + // Read required MaskName parameter from configuration + std::string MaskFieldName; + Error Err = Options.get("MaskName", MaskFieldName); + + if (Err.isFail()) { + ABORT_ERROR("ExtractRegion Op: Required parameter 'MaskName' not " + "found in configuration"); + } + + // Retrieve the mask field from registry + auto MaskField = Field::get(MaskFieldName); + if (MaskField == nullptr) { + ABORT_ERROR("ExtractRegionOp: Mask field '{}' not found in Field " + "registry", + MaskFieldName); + } + + // Verify mask is 1D integer array + if (MaskField->getNumDims() != 1) { + ABORT_ERROR("ExtractRegionOp: Mask field '{}' must be 1D, got {}D", + MaskFieldName, MaskField->getNumDims()); + } + + // Extract mask data + RegionalMask = MaskField->template getDataArray(); + + // Retrieve input Field to get dimensions and metadata + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + // Support 1D (horizontal), 2D (horizontal × vertical), and 3D arrays + constexpr I4 InputRank = ArrayT::rank; + + // Get dimension names from input field + std::vector InputDimNames; + InputField->getDimNames(InputDimNames); + + // Get dimension name from mask field + std::vector MaskDimNames; + MaskField->getDimNames(MaskDimNames); + + // Determine the horizontal dimension name based on rank + std::string InputHorizDimName; + if constexpr (InputRank == 1) { + // 1D: horizontal only + InputHorizDimName = InputDimNames[0]; + NHorizDim = InputData.extent(0); + NVertSize = 1; // No vertical dimension + } else if constexpr (InputRank == 2) { + // 2D: NHorizDim × NVertSize (horizontal is first dimension) + InputHorizDimName = InputDimNames[0]; + NHorizDim = InputData.extent(0); + NVertSize = InputData.extent(1); + } else if constexpr (InputRank == 3) { + // 3D: Dim0 × NHorizDim × NVertSize (horizontal is second dimension) + InputHorizDimName = InputDimNames[1]; + Dim0 = InputData.extent(0); + NHorizDim = InputData.extent(1); + NVertSize = InputData.extent(2); + } else { + ABORT_ERROR("ExtractRegionOp: Unsupported rank {}. " + "Supports 1D, 2D, and 3D arrays", + InputRank); + } + + // Verify mask dimension name matches input horizontal dimension name + if (MaskDimNames[0] != InputHorizDimName) { + ABORT_ERROR("ExtractRegionOp: Mask dimension name '{}' does not match " + "input horizontal dimension name '{}'", + MaskDimNames[0], InputHorizDimName); + } + + // Construct output field name: InputName + "_" + RegionName + // Strip "Mask" suffix from MaskFieldName if present for cleaner names + std::string RegionName = MaskFieldName; + if (RegionName.size() > 4 && + RegionName.substr(RegionName.size() - 4) == "Mask") { + RegionName = RegionName.substr(0, RegionName.size() - 4); + } + + std::string OutputFieldName = InputNames[0] + "_" + RegionName; + OutputNames = {OutputFieldName}; + InstanceName = OutputFieldName; + + // Get dimension info from input + auto NDims = InputField->getNumDims(); + std::vector DimNames; + InputField->getDimNames(DimNames); + + // Get input metadata + std::string InputDescr, InputUnits, InputStdName; + ScalarT InputValidMin, InputValidMax; + + InputField->getMetadata("Description", InputDescr); + InputField->getMetadata("Units", InputUnits); + InputField->getMetadata("StdName", InputStdName); + InputField->getMetadata("ValidMin", InputValidMin); + InputField->getMetadata("ValidMax", InputValidMax); + + // Create output Field with same dimensions as input + auto OutputField = Field::create( + OutputNames[0], + InputDescr + " (region: " + RegionName + ")", // Description + InputUnits, // Units (unchanged) + InputStdName, // Standard name (unchanged) + InputValidMin, // Min valid (unchanged) + InputValidMax, // Max valid (unchanged) + NDims, // Rank + DimNames // Dimension names + ); + + // Allocate output data array matching input layout + OutputData = OutputArrayT(OutputNames[0] + "_out", InputData.layout()); + + // Attach output data array to Field + OutputField->template attachData(OutputData); + + // Attach the regional mask to the output Field + // If input already has a mask, compute intersection + if (InputField->hasRegionalMask()) { + Array1DI4 InputMask = InputField->getRegionalMask(); + Array1DI4 IntersectionMask = + Array1DI4("IntersectionMask", RegionalMask.extent(0)); + + // Compute intersection: both masks must be 1 for result to be 1 + auto LocalInputMask = InputMask; + auto LocalRegionalMask = RegionalMask; + auto LocalIntersectionMask = IntersectionMask; + parallelFor( + {static_cast(RegionalMask.extent(0))}, KOKKOS_LAMBDA(int I) { + LocalIntersectionMask(I) = + LocalInputMask(I) * LocalRegionalMask(I); + }); + + OutputField->setRegionalMask(IntersectionMask); + } else { + // No existing mask, just attach this one + OutputField->setRegionalMask(RegionalMask); + } + + } // end constructor + + /// Initializes the operator after all Fields exist. Determines the index + /// space (cells, edges, or vertices) from the input field's horizontal + /// dimension name and stores the appropriate MinLayer/MaxLayer arrays + /// from VertCoord for bounding the inner vertical loop. + void initialize(const MachEnv *Env, const HorzMesh *InMesh, + const VertCoord *InVCoord, Config Options) override { + + AnalysisOperator::initialize(Env, InMesh, InVCoord, Options); + + constexpr I4 InputRank = ArrayT::rank; + + // For 2D/3D arrays, get MinLayer/MaxLayer for partial columns + if constexpr (InputRank > 1) { + auto Field1 = Field::get(InputNames[0]); + std::vector DimNames; + Field1->getDimNames(DimNames); + // Horizontal dimension is 2nd-to-last for 2D/3D + std::string IndexSpaceName = DimNames[InputRank - 2]; + + if (IndexSpaceName == "NCells") { + MinLayer = VCoord->MinLayerCell; + MaxLayer = VCoord->MaxLayerCell; + NHorizOwned = Mesh->NCellsOwned; + } else if (IndexSpaceName == "NEdges") { + MinLayer = VCoord->MinLayerEdgeBot; + MaxLayer = VCoord->MaxLayerEdgeTop; + NHorizOwned = Mesh->NEdgesOwned; + } else if (IndexSpaceName == "NVertices") { + MinLayer = VCoord->MinLayerVertexBot; + MaxLayer = VCoord->MaxLayerVertexTop; + NHorizOwned = Mesh->NVerticesOwned; + } else { + ABORT_ERROR("ExtractRegionOp: Unknown index space {}", + IndexSpaceName); + } + } else { + // For 1D arrays, determine owned count + auto Field1 = Field::get(InputNames[0]); + std::vector DimNames; + Field1->getDimNames(DimNames); + std::string IndexSpaceName = DimNames[0]; + + if (IndexSpaceName == "NCells") { + NHorizOwned = Mesh->NCellsOwned; + } else if (IndexSpaceName == "NEdges") { + NHorizOwned = Mesh->NEdgesOwned; + } else if (IndexSpaceName == "NVertices") { + NHorizOwned = Mesh->NVerticesOwned; + } else { + ABORT_ERROR("ExtractRegionOp: Unknown index space {}", + IndexSpaceName); + } + } + + } // end initialize + + /// Computes the regional extraction by multiplying input data with the + /// regional mask. Uses hierarchical parallelism with MinLayer/MaxLayer + /// bounding for partial columns. The mask (1D horizontal) is applied + /// across all vertical layers and extra dimensions. + void compute(const TimeInstant &TimeStamp ///< [in] current timestamp + ) override { + + // Get array rank for conditional logic + constexpr I4 InputRank = ArrayT::rank; + + // Retrieve input Field and extract data array + auto InputField = Field::get(InputNames[0]); + auto InputData = InputField->template getDataArray(); + + OMEGA_SCOPE(LocOutput, OutputData); + OMEGA_SCOPE(LocMask, RegionalMask); + + if constexpr (InputRank == 1) { + // 1D case: horizontal field only + parallelFor( + {NHorizOwned}, KOKKOS_LAMBDA(int iHoriz) { + const Real MaskVal = static_cast(LocMask(iHoriz)); + if (MaskVal != 0) { + LocOutput(iHoriz) = static_cast( + static_cast(InputData(iHoriz)) * + static_cast(MaskVal)); + } + }); + + } else if constexpr (InputRank == 2) { + // 2D case: hierarchical parallelism over horizontal × vertical + // Inner loop bounded by MinLayer/MaxLayer for partial columns + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + + parallelForOuter( + "ExtractRegion2D", LaunchConfig({NHorizOwned}), + KOKKOS_LAMBDA(int iHoriz, const TeamMember &Team) { + const Real MaskVal = static_cast(LocMask(iHoriz)); + if (MaskVal != 0) { + const I4 KMin = LocMinLayer(iHoriz); + const I4 KMax = LocMaxLayer(iHoriz); + const I4 KRange = vertRange(KMin, KMax); + + parallelForInner( + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + LocOutput(iHoriz, K) = static_cast( + static_cast(InputData(iHoriz, K)) * + MaskVal); + }); + } + }); + + } else if constexpr (InputRank == 3) { + // 3D case: hierarchical parallelism over dim0 × horizontal + // Inner loop bounded by MinLayer/MaxLayer for partial columns + OMEGA_SCOPE(LocMinLayer, MinLayer); + OMEGA_SCOPE(LocMaxLayer, MaxLayer); + OMEGA_SCOPE(LocDim0, Dim0); + + parallelForOuter( + "ExtractRegion3D", LaunchConfig({LocDim0, NHorizOwned}), + KOKKOS_LAMBDA(int i0, int iHoriz, const TeamMember &Team) { + const Real MaskVal = static_cast(LocMask(iHoriz)); + if (MaskVal != 0) { + const I4 KMin = LocMinLayer(iHoriz); + const I4 KMax = LocMaxLayer(iHoriz); + const I4 KRange = vertRange(KMin, KMax); + + parallelForInner( + Team, KRange, INNER_LAMBDA(int KIdx) { + const I4 K = KMin + KIdx; + LocOutput(i0, iHoriz, K) = static_cast( + static_cast(InputData(i0, iHoriz, K)) * + MaskVal); + }); + } + }); + } + + // Update cache validity markers + LastComputed = TimeStamp; + FieldComputed = true; + + } // end compute + + private: + /// Output data array holding the masked field values + OutputArrayT OutputData; + + /// The regional mask (1D horizontal) + Array1DI4 RegionalMask; + + /// Number of points in horizontal dimension + I4 NHorizDim; + + /// Number of owned points in horizontal dimension (for MPI correctness) + I4 NHorizOwned; + + /// Vertical size of the field array + I4 NVertSize; + + /// Size of first dimension for 3D arrays + I4 Dim0; + + /// Min active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MinLayer; + + /// Max active layer index for each horizontal point (only for 2D/3D arrays) + Array1DI4 MaxLayer; + +}; // end class ExtractRegionOp + +} // end namespace OMEGA + +#endif From 377432c9f490bec92d0f7445af9e03ab77b7a6cd Mon Sep 17 00:00:00 2001 From: Brian O'Neill Date: Mon, 27 Jul 2026 00:17:23 -0400 Subject: [PATCH 09/35] Add register1D/2D/2DReal variant methods to AnalysisOpFactory --- .../omega/src/analysis/AnalysisOpFactory.h | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) diff --git a/components/omega/src/analysis/AnalysisOpFactory.h b/components/omega/src/analysis/AnalysisOpFactory.h index f394f8118e23..c58e5708b34f 100644 --- a/components/omega/src/analysis/AnalysisOpFactory.h +++ b/components/omega/src/analysis/AnalysisOpFactory.h @@ -155,6 +155,134 @@ class AnalysisOpFactory { #undef REGISTER_VARIANT } // end registerAllArrayVariants + /// Registers only 1D array type variants of a templated operator class. + /// Used for operators that only support 1D arrays. Uses SFINAE to prevent + /// template instantiation for non-1D array types at compile time. + + // Helper: enabled only for 1D arrays - registers the operator variant + template