Simplify Hamiltonian Creation (#133)

* now create hamiltonian with more tensor-y operations rather than in for loop, hopefully saving some time
Also requires some fiddling with shape of different internal tensor values

Author: ewanwm

nuTens/propagator/base-matter-solver.hpp

@@ -35,7 +35,8 @@ class BaseMatterSolver
       NT_PROFILE();
 
       energies = newEnergies;
-      energiesRed = energies.getValues({"...", 0});
+      energiesRed = energies.getValues({"..."});
+      energiesRed.unsqueeze(-1);
 
       hamiltonian = Tensor::zeros({energies.getBatchDim(), nGenerations, nGenerations}, dtypes::kComplexFloat).requiresGrad(false);
     }

nuTens/propagator/const-density-solver.cpp

@@ -5,18 +5,11 @@ using namespace nuTens;
 void ConstDensityMatterSolver::calculateEigenvalues(Tensor &eigenvectors, Tensor &eigenvalues)
 {
     NT_PROFILE();
-    
-    for (int i = 0; i < nGenerations; i++)
-    {
-        for (int j = 0; j < nGenerations; j++)
-        {
-            hamiltonian.setValue({"...", i, j},
-                                 Tensor::div(diagMassMatrix.getValues({i, j}), energiesRed) -
-                                     electronOuter.getValues({i, j}));
-        }
-    }
+
+    hamiltonian.setValue({"..."}, (Tensor::div(diagMassMatrix, energiesRed) - electronOuter));
 
     Tensor::eigh(hamiltonian, eigenvalues, eigenvectors);
+    
 }
 
 void ConstDensityMatterSolver::buildElectronOuterProduct() 
@@ -29,7 +22,6 @@ void ConstDensityMatterSolver::buildElectronOuterProduct()
         electronOuter =
             Tensor::scale(Tensor::outer(mixingMatrix.getValues({0, 0, "..."}).conj(), mixingMatrix.getValues({0, 0, "..."})),
                           -nuTens::constants::Groot2 * density);
-        
     }
 
     else 
@@ -38,4 +30,6 @@ void ConstDensityMatterSolver::buildElectronOuterProduct()
             Tensor::scale(Tensor::outer(mixingMatrix.getValues({0, 0, "..."}), mixingMatrix.getValues({0, 0, "..."}).conj()),
                           nuTens::constants::Groot2 * density);
     }
+
+    electronOuter.unsqueeze(0);
 }

nuTens/propagator/const-density-solver.hpp

@@ -84,7 +84,7 @@ class ConstDensityMatterSolver : public BaseMatterSolver
         Tensor diag = Tensor::scale(Tensor::mul(m, m), 0.5);
 
         // construct the diagonal mass^2 matrix used in the hamiltonian
-        diagMassMatrix = Tensor::diag(diag).requiresGrad(false);
+        diagMassMatrix = Tensor::diag(diag).requiresGrad(false).unsqueeze(0);
     }