field_state.hpp

Go to the documentation of this file.

// Copyright (c) Lawrence Livermore National Security, LLC and
// other Smith Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include "gretl/data_store.hpp"
#include "gretl/state.hpp"
#include "gretl/create_state.hpp"
#include "smith/physics/state/state_manager.hpp"
 
namespace smith {
 
using FEFieldPtr = std::shared_ptr<FiniteElementState>;                               
using FEDualPtr = std::shared_ptr<FiniteElementDual>;                                 
using FieldState = gretl::State<FEFieldPtr, FEDualPtr>;                               
using ReactionState = gretl::State<FEDualPtr, FEFieldPtr>;                            
using FieldVecState = gretl::State<std::vector<FEFieldPtr>, std::vector<FEDualPtr>>;  
using DoubleState = gretl::State<double, double>;                                     
 
struct zero_dual_from_state {
  auto operator()(const smith::FEFieldPtr& f) const
  {
    return std::make_shared<smith::FiniteElementDual>(f->space(), f->name() + "_dual");
  };
};
 
struct zero_state_from_dual {
  auto operator()(const smith::FEDualPtr& f) const
  {
    return std::make_shared<smith::FiniteElementState>(f->space(), f->name() + "_undual");
  };
};
 
inline FieldState createFieldState(gretl::DataStore& dataStore, const smith::FEFieldPtr& s)
{
  return dataStore.create_state<smith::FEFieldPtr, smith::FEDualPtr>(s, zero_dual_from_state());
}
 
template <typename function_space>
FieldState createFieldState(gretl::DataStore& dataStore, function_space space, const std::string& name,
                            const std::string& mesh_tag)
{
  auto f = std::make_shared<FiniteElementState>(StateManager::newState(space, name, mesh_tag));
  return createFieldState(dataStore, f);
}
 
inline ReactionState createReactionState(gretl::DataStore& dataStore, const smith::FEDualPtr& s)
{
  return dataStore.create_state<smith::FEDualPtr, smith::FEFieldPtr>(s, zero_state_from_dual());
}
 
template <typename function_space>
ReactionState createReactionState(gretl::DataStore& dataStore, function_space space, const std::string& name,
                                  const std::string& mesh_tag)
{
  auto f = std::make_shared<FiniteElementDual>(StateManager::newDual(space, name, mesh_tag));
  return createReactionState(dataStore, f);
}
 
FieldState square(const FieldState& state);
 
gretl::State<double> innerProduct(const FieldState& a, const FieldState& b);
 
gretl::State<double> innerProduct(const ReactionState& a, const ReactionState& b);
 
FieldState axpby(double a, const FieldState& x, double b, const FieldState& y);
 
FieldState zeroCopy(const FieldState& x);
 
inline FieldState weighted_average(const FieldState& a, const FieldState& b, double weight)
{
  return axpby(weight, a, 1.0 - weight, b);
}
 
FieldState axpby(const gretl::State<double>& a, const FieldState& x, const gretl::State<double>& b,
                 const FieldState& y);
 
struct FieldStateWeightedSum {
  FieldStateWeightedSum(const std::vector<double>& w, const std::vector<FieldState>& f)
      : weights_(w), weighted_fields_(f)
  {
  }
 
  FieldStateWeightedSum(const std::vector<gretl::State<double>>& w, const std::vector<FieldState>& f,
                        double initial_scaling)
      : differentiable_weights_(w),
        differentiably_weighted_fields_(f),
        differentiable_scale_factors_(w.size(), initial_scaling)
  {
  }
 
  FieldStateWeightedSum(const FieldStateWeightedSum& old) = default;
 
  FieldStateWeightedSum& operator=(const FieldStateWeightedSum& old) = default;
 
  FieldStateWeightedSum& operator+=(const FieldStateWeightedSum& b);
 
  FieldStateWeightedSum& operator-=(const FieldStateWeightedSum& b);
 
  FieldStateWeightedSum& operator*=(double weight);
 
  FieldStateWeightedSum operator-() const;
 
  std::vector<double> weights_;                               
  std::vector<FieldState> weighted_fields_;                   
  std::vector<gretl::State<double>> differentiable_weights_;  
  std::vector<FieldState> differentiably_weighted_fields_;    
  std::vector<double> differentiable_scale_factors_;          
 
  operator FieldState() const;
};
 
FieldStateWeightedSum operator*(double a, const FieldState& b);
 
FieldStateWeightedSum operator*(const FieldState& b, double a);
 
FieldStateWeightedSum operator*(double a, const FieldStateWeightedSum& b);
 
FieldStateWeightedSum operator*(const FieldStateWeightedSum& b, double a);
 
FieldStateWeightedSum operator*(const gretl::State<double>& a, const FieldState& b);
 
FieldStateWeightedSum operator*(const FieldState& b, const gretl::State<double>& a);
 
FieldStateWeightedSum operator+(const FieldState& x, const FieldState& y);
 
FieldStateWeightedSum operator-(const FieldState& x, const FieldState& y);
 
FieldStateWeightedSum operator+(const FieldStateWeightedSum& ax, const FieldStateWeightedSum& by);
 
FieldStateWeightedSum operator-(const FieldStateWeightedSum& ax, const FieldStateWeightedSum& by);
 
FieldStateWeightedSum operator+(const FieldStateWeightedSum& ax, const FieldState& y);
 
FieldStateWeightedSum operator+(const FieldState& y, const FieldStateWeightedSum& ax);
 
FieldStateWeightedSum operator-(const FieldStateWeightedSum& ax, const FieldState& by);
 
FieldStateWeightedSum operator-(const FieldState& ax, const FieldStateWeightedSum& by);
 
// TODO:
// Add multplication of WeightedSum by a differentiable State<double> for improve efficiency
// Consider adding divide operators, maybe component-wise things as well
 
// Utilty functions for easily getting spaces from FieldStates
 
inline mfem::ParFiniteElementSpace& space(FieldState field) { return field.get()->space(); }
 
inline std::vector<const mfem::ParFiniteElementSpace*> spaces(const std::vector<FieldState>& states,
                                                              const std::vector<FieldState>& params = {})
{
  std::vector<const mfem::ParFiniteElementSpace*> spaces;
  for (auto s : states) {
    spaces.push_back(&s.get()->space());
  }
  for (auto s : params) {
    spaces.push_back(&s.get()->space());
  }
  return spaces;
};
 
inline std::vector<FiniteElementState*> getFieldPointers(std::vector<FieldState>& states,
                                                         std::vector<FieldState> params = {})
{
  std::vector<FiniteElementState*> pointers;
  for (auto& t : states) {
    pointers.push_back(t.get().get());
  }
  for (auto& p : params) {
    pointers.push_back(p.get().get());
  }
  return pointers;
}
 
inline std::vector<const FiniteElementState*> getConstFieldPointers(const std::vector<FieldState>& states,
                                                                    const std::vector<FieldState>& params = {})
{
  std::vector<const FiniteElementState*> pointers;
  for (auto& t : states) {
    pointers.push_back(t.get().get());
  }
  for (auto& p : params) {
    pointers.push_back(p.get().get());
  }
  return pointers;
}
 
}  // namespace smith