tensor_network.hpp

 
#ifndef EXATN_NUMERICS_TENSOR_NETWORK_HPP_
#define EXATN_NUMERICS_TENSOR_NETWORK_HPP_
 
#include "tensor_basic.hpp"
#include "tensor_connected.hpp"
#include "tensor_op_factory.hpp"
#include "network_build_factory.hpp"
#include "contraction_seq_optimizer.hpp"
 
#include <unordered_map>
#include <map>
#include <vector>
#include <list>
#include <string>
#include <memory>
 
namespace exatn{
 
namespace numerics{
 
class TensorNetwork{
public:
 
 using Iterator = typename std::unordered_map<unsigned int, TensorConn>::iterator; //iterator
 using ConstIterator = typename std::unordered_map<unsigned int, TensorConn>::const_iterator; //constant iterator
 
 TensorNetwork();
 TensorNetwork(const std::string & name);
 TensorNetwork(const std::string & name,                    //in: tensor network name
               std::shared_ptr<Tensor> output_tensor,       //in: fully specified output tensor of the tensor network
               const std::vector<TensorLeg> & output_legs); //in: fully specified output tensor legs
 TensorNetwork(const std::string & name,                    //in: tensor network name
               const std::string & tensor_network,          //in: tensor network expression (symbolic math expression)
               const std::map<std::string,std::shared_ptr<Tensor>> & tensors); //in: participating tensors identified by their names
 TensorNetwork(const std::string & name,                    //in: tensor network name
               std::shared_ptr<Tensor> output_tensor,       //in: output tensor of the tensor network
               NetworkBuilder & builder);                   //in: specific tensor network builder
 
 TensorNetwork(const TensorNetwork &) = default;
 TensorNetwork & operator=(const TensorNetwork &) = default;
 TensorNetwork(TensorNetwork &&) noexcept = default;
 TensorNetwork & operator=(TensorNetwork &&) noexcept = default;
 virtual ~TensorNetwork() = default;
 
 void printIt() const;
 
 bool isEmpty() const;
 
 bool isExplicit() const;
 
 bool isFinalized() const;
 
 unsigned int getRank() const;
 
 unsigned int getNumTensors() const;
 
 unsigned int getMaxTensorId() const;
 
 const std::string & getName() const;
 
 void rename(const std::string & name);
 
 std::shared_ptr<Tensor> getTensor(unsigned int tensor_id,
                                   bool * conjugated = nullptr);
 
 const std::vector<TensorLeg> * getTensorConnections(unsigned int tensor_id);
 
 inline Iterator begin() {return tensors_.begin();}
 inline Iterator end() {return tensors_.end();}
 inline ConstIterator cbegin() const {return tensors_.cbegin();}
 inline ConstIterator cend() const {return tensors_.cend();}
 
 bool finalize(bool check_validity = false);
 
 bool appendTensor(unsigned int tensor_id,                     //in: appended tensor id (unique within the tensor network)
                   std::shared_ptr<Tensor> tensor,             //in: appended tensor
                   const std::vector<TensorLeg> & connections, //in: tensor connections (fully specified)
                   bool conjugated = false,                    //in: complex conjugation flag for the appended tensor
                   bool leg_matching_check = true);            //in: tensor leg matching check
 
 bool appendTensor(unsigned int tensor_id,                                                  //in: appended tensor id (unique within the tensor network)
                   std::shared_ptr<Tensor> tensor,                                          //in: appended tensor
                   const std::vector<std::pair<unsigned int, unsigned int>> & pairing,      //in: leg pairing: output tensor mode -> appended tensor mode
                   const std::vector<LegDirection> & leg_dir = std::vector<LegDirection>{}, //in: optional leg directions (for all tensor modes)
                   bool conjugated = false);                                                //in: complex conjugation flag for the appended tensor
 
 bool appendTensorGate(unsigned int tensor_id,                    //in: appended tensor id (unique within the tensor network)
                       std::shared_ptr<Tensor> tensor,            //in: appended tensor gate (operator)
                       const std::vector<unsigned int> & pairing, //in: leg pairing: output tensor modes (half-rank)
                       bool conjugated = false);                  //in: complex conjugation flag for the appended tensor gate
 
 bool appendTensorNetwork(TensorNetwork && network,                                            //in: appended tensor network
                          const std::vector<std::pair<unsigned int, unsigned int>> & pairing); //in: leg pairing: output tensor mode (primary) -> output tensor mode (appended)
 
 bool appendTensorNetworkGate(TensorNetwork && network,                   //in: appended tensor network gate (operator)
                              const std::vector<unsigned int> & pairing); //in: leg pairing: output tensor modes of the primary network (half-rank)
 
 bool reorderOutputModes(const std::vector<unsigned int> & order); //in: new order of the output tensor modes (N2O)
 
 bool deleteTensor(unsigned int tensor_id); //in: id of the tensor to be deleted
 
 bool mergeTensors(unsigned int left_id,   //in: left tensor id (present in the tensor network)
                   unsigned int right_id,  //in: right tensor id (present in the tensor network)
                   unsigned int result_id, //in: result tensor id (absent in the tensor network, to be appended)
                   std::string * contr_pattern = nullptr); //inout: corresponding tensor contraction pattern (owned by the caller)
 
 bool splitTensor(unsigned int tensor_id,                //in: id of the tensor to be split into two tensors
                  unsigned int left_tensor_id,           //in: id of the left tensor obtained via splitting
                  const std::string & left_tensor_name,  //in: name of the left tensor
                  unsigned int right_tensor_id,          //in: id of the right tensor obtained via splitting
                  const std::string & right_tensor_name, //in: name of the right tensor
                  const TensorShape & contracted_dims,   //in: dimension extents of the contracted (new) dimensions connecting two tensors after splitting
                  const std::vector<int> & right_dims);  //in: assignment of original tensor dimensions to new tensors (0: left, 1: right tensor)
 
 void conjugate();
 
 double getContractionCost(unsigned int left_id,  //in: left tensor id (present in the tensor network)
                           unsigned int right_id, //in: right tensor id (present in the tensor network)
                           double * arithm_intensity = nullptr, //out: arithmetic intensity of the tensor contraction
                           bool adjust_cost = false); //in: whether or not to adjust the flops cost due to arithmetic intensity
 
 std::list<std::shared_ptr<TensorOperation>> & getOperationList(const std::string & contr_seq_opt_name = "dummy");
 
protected:
 
 TensorConn * getTensorConn(unsigned int tensor_id);
 
 std::vector<TensorConn*> getTensorConnAll();
 
 bool checkConnections(unsigned int tensor_id);
 bool checkConnections();
 
 void updateConnections(unsigned int tensor_id); //in: id of the tensor whose connections were modified
 
 void updateConnectionsFromInputTensors();
 
 void invalidateContractionSequence();
 
 double determineContractionSequence(ContractionSeqOptimizer & contr_seq_optimizer);
 
private:
 
 int explicit_output_;                                  //whether or not the output tensor has been fully specified during construction
 int finalized_;                                        //finalization status of the tensor network
 std::string name_;                                     //tensor network name
 std::unordered_map<unsigned int, TensorConn> tensors_; //tensors connected to each other via legs (tensor connections)
                                                        //map: Non-negative tensor id --> Connected tensor
 double contraction_seq_flops_; //flop estimate for the determined tensor contraction sequence
 std::list<ContrTriple> contraction_seq_; //cached tensor contraction sequence
 std::list<std::shared_ptr<TensorOperation>> operations_; //cached tensor operations required for evaluating the tensor network
};
 
} //namespace numerics
 
template<typename... Args>
inline std::shared_ptr<numerics::TensorNetwork> makeSharedTensorNetwork(Args&&... args)
{
 return std::make_shared<numerics::TensorNetwork>(args...);
}
 
} //namespace exatn
 
#endif //EXATN_NUMERICS_TENSOR_NETWORK_HPP_