From 5167ff3a4dbb1d7ea16c5bfb4e94783e6b0cdb6b Mon Sep 17 00:00:00 2001
From: Joachim Schoeberl <joachim.schoeberl@tuwien.ac.at>
Date: Wed, 14 Jul 2021 16:53:41 +0200
Subject: [PATCH] trace with averaging

---
 comp/fespace.hpp     |  2 +-
 comp/l2hofespace.cpp | 26 +++++++++++++++++++++++++-
 comp/l2hofespace.hpp |  2 +-
 comp/python_comp.cpp |  2 +-
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/comp/fespace.hpp b/comp/fespace.hpp
index e4b3a3638..d0fe52101 100644
--- a/comp/fespace.hpp
+++ b/comp/fespace.hpp
@@ -698,7 +698,7 @@ ANY                  1 1 1 1 | 15
     virtual void ApplyM(CoefficientFunction * rho, BaseVector & vec, Region * definedon,
                         LocalHeap & lh) const;
 
-    virtual shared_ptr<BaseMatrix> GetTraceOperator (shared_ptr<FESpace> tracespace) const
+    virtual shared_ptr<BaseMatrix> GetTraceOperator (shared_ptr<FESpace> tracespace, bool avg) const
     { throw Exception("GetTraceOperator not overloaded"); }
 
     virtual shared_ptr<BaseMatrix> ConvertL2Operator (shared_ptr<FESpace> l2space) const;
diff --git a/comp/l2hofespace.cpp b/comp/l2hofespace.cpp
index 4d68c8b53..2a97fdcb5 100644
--- a/comp/l2hofespace.cpp
+++ b/comp/l2hofespace.cpp
@@ -1147,7 +1147,7 @@ global system.
 
 
   shared_ptr<BaseMatrix> L2HighOrderFESpace ::
-  GetTraceOperator (shared_ptr<FESpace> tracespace) const
+  GetTraceOperator (shared_ptr<FESpace> tracespace, bool avg) const
   {
     LocalHeap lh(1000000);
     Array<short> classnr(ma->GetNE());
@@ -1170,6 +1170,14 @@ global system.
 
     // size_t ne = ma->GetNE();
 
+    shared_ptr<VVector<double>> cnt;
+    if (avg)
+      {
+        cnt = make_shared<VVector<double>>(tracespace->GetNDof());
+        *cnt = 0;
+      }
+
+    
     for (auto elclass_inds : table)
       {
         if (elclass_inds.Size() == 0) continue;
@@ -1192,6 +1200,10 @@ global system.
             tracespace->GetDofNrs(ei, dnumsy);
             xdofs[i] = dnumsx;
             ydofs[i] = dnumsy;
+            
+            if (avg)
+              for (auto d : dnumsy)
+                (*cnt)(d) += 1;
           }
 
         auto mat = make_shared<ConstantElementByElementMatrix>
@@ -1203,6 +1215,18 @@ global system.
         else
           sum = mat;
       }
+
+    if (avg)
+      {
+        for (size_t i : Range(cnt->Size()))
+          if ( (*cnt)(i) != 0)
+            (*cnt)(i) = 1 / (*cnt)(i);
+
+        auto diag = make_shared<DiagonalMatrix<double>> (cnt);
+        sum = make_shared<ProductMatrix> (diag, sum);
+      }
+
+      
     return sum;
   }
 
diff --git a/comp/l2hofespace.hpp b/comp/l2hofespace.hpp
index a5ed01033..aa00916bd 100644
--- a/comp/l2hofespace.hpp
+++ b/comp/l2hofespace.hpp
@@ -112,7 +112,7 @@ namespace ngcomp
     virtual void ApplyM (CoefficientFunction * rho, BaseVector & vec, Region * definedon,
                          LocalHeap & lh) const override;
 
-    virtual shared_ptr<BaseMatrix> GetTraceOperator (shared_ptr<FESpace> tracespace) const override;
+    virtual shared_ptr<BaseMatrix> GetTraceOperator (shared_ptr<FESpace> tracespace, bool avg) const override;
     virtual void GetTrace (const FESpace & tracespace, const BaseVector & in, BaseVector & out, bool avg,
                            LocalHeap & lh) const override;
     virtual void GetTraceTrans (const FESpace & tracespace, const BaseVector & in, BaseVector & out, bool avg,
diff --git a/comp/python_comp.cpp b/comp/python_comp.cpp
index c9e1bf2d2..84d18e862 100644
--- a/comp/python_comp.cpp
+++ b/comp/python_comp.cpp
@@ -1064,7 +1064,7 @@ rho : ngsolve.fem.CoefficientFunction
     .def ("TraceOperator", [] (shared_ptr<FESpace> self, shared_ptr<FESpace> tracespace,
                                bool avg) -> shared_ptr<BaseMatrix>
           {
-            return self->GetTraceOperator(tracespace);
+            return self->GetTraceOperator(tracespace, avg);
             // return make_shared<ApplyTrace> (self, tracespace, avg, glh);             
           }, py::arg("tracespace"), py::arg("average"))
     .def ("ConvertL2Operator", [] (shared_ptr<FESpace> self, shared_ptr<FESpace> l2space)