Add all the Basic Binary Operators

Author: Neel-Shah-29

tmva/sofie/inc/TMVA/ROperator_Add.hxx

@@ -1,5 +1,5 @@
-#ifndef TMVA_SOFIE_ROPERATOR_ADD
-#define TMVA_SOFIE_ROPERATOR_ADD
+#ifndef TMVA_SOFIE_ROPERATOR_BinaryOperator
+#define TMVA_SOFIE_ROPERATOR_BinaryOperator
 
 #include "TMVA/SOFIE_common.hxx"
 #include "TMVA/ROperator.hxx"
@@ -11,8 +11,9 @@ namespace TMVA{
 namespace Experimental{
 namespace SOFIE{
 
+enum RModel_BasicBinaryOperator { Add, Sub, Mul, Div };
 template <typename T>
-class ROperator_Add final : public ROperator
+class ROperator_BinaryOperator final : public ROperator
 {
 
 private:
@@ -23,8 +24,17 @@ private:
    std::vector<size_t> fShape;
 
 public:
-   ROperator_Add(){}
-   ROperator_Add(std::string nameX1, std::string nameX2, std::string nameY):
+
+      std::string Name() {
+      if (RModel_BasicBinaryOperator == Add)  return "Add";
+      else if (RModel_BasicBinaryOperator == Sub)  return "Sub";
+      else if (RModel_BasicBinaryOperator == Mul ) return "Mul";
+      else if (RModel_BasicBinaryOperator == Div ) return "Div";
+      return "Invalid";
+   }
+
+   ROperator_BinaryOperator(){}
+   ROperator_BinaryOperator(std::string nameX1, std::string nameX2, std::string nameY):
       fNX1(UTILITY::Clean_name(nameX1)), fNX2(UTILITY::Clean_name(nameX2)), fNY(UTILITY::Clean_name(nameY)){}
 
    // type of output given input 
@@ -42,16 +52,16 @@ public:
    void Initialize(RModel& model){
       // input must be a graph input, or already initialized intermediate tensor
       if (model.CheckIfTensorAlreadyExist(fNX1) == false){
-         throw std::runtime_error(std::string("TMVA SOFIE Add Op Input Tensor ") + fNX1 + "is not found in model");
+         throw std::runtime_error(std::string("TMVA SOFIE Binary Op Input Tensor ") + fNX1 + "is not found in model");
       }
       if (model.CheckIfTensorAlreadyExist(fNX2) == false) {
-         throw std::runtime_error(std::string("TMVA SOFIE Add Op Input Tensor ") + fNX1 + "is not found in model");
+         throw std::runtime_error(std::string("TMVA SOFIE Binary Op Input Tensor ") + fNX2 + "is not found in model");
       }
       auto shapeX1 = model.GetTensorShape(fNX1);
       auto shapeX2 = model.GetTensorShape(fNX2);
       // assume same shape X1 and X2 
       if (shapeX1 != shapeX2) {
-         std::string msg = "TMVA SOFIE Add Op: Support only inputs with same shape, shape 1 is " +
+         std::string msg = "TMVA SOFIE Binary Op: Support only inputs with same shape, shape 1 is " +
                            ConvertShapeToString(shapeX1) + "shape 2 is " + ConvertShapeToString(shapeX2);
          throw std::runtime_error(msg);
       }
@@ -63,18 +73,36 @@ public:
    std::string Generate(std::string OpName){
       OpName = "op_" + OpName;
       if (fShape.empty()) {
-         throw std::runtime_error("TMVA SOFIE Add called to Generate without being initialized first");
+         throw std::runtime_error("TMVA SOFIE binary operator called to Generate without being initialized first");
       }
       std::stringstream out;
       // int length = 1;
       // for(auto& i: fShape){
       //    length *= i;
       // }
+      
       size_t length = ConvertShapeToLength(fShape);
-      out << "\n//------ Add\n";
+      out << "\n//----  operator " << Name() << "  " << OpName << "\n";
+      if(RModel_BasicBinaryOperator == Add){
       out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
       out << SP << SP << "tensor_" << fNY << "[id] = tensor_" << fNX1 << "[id] + tensor_" << fNX2 << "[id];\n";
       out << SP << "}\n";
+      }
+      if(RModel_BasicBinaryOperator == Sub){
+      out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
+      out << SP << SP << "tensor_" << fNY << "[id] = tensor_" << fNX1 << "[id] - tensor_" << fNX2 << "[id];\n";
+      out << SP << "}\n";
+      }
+      if(RModel_BasicBinaryOperator == Mul){
+      out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
+      out << SP << SP << "tensor_" << fNY << "[id] = tensor_" << fNX1 << "[id] * tensor_" << fNX2 << "[id];\n";
+      out << SP << "}\n";
+      }
+      if(RModel_BasicBinaryOperator == Div){
+      out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
+      out << SP << SP << "tensor_" << fNY << "[id] = tensor_" << fNX1 << "[id] / tensor_" << fNX2 << "[id];\n";
+      out << SP << "}\n";
+      }
       return out.str();
    }
 
@@ -85,4 +113,4 @@ public:
 }//TMVA
 
 
-#endif //TMVA_SOFIE_ROPERATOR_Add
+#endif //TMVA_SOFIE_ROPERATOR_BinaryOperator