diff --git a/hw3-heat-fft/compute_temperature.cc b/hw3-heat-fft/compute_temperature.cc
index dda4b0b4..62419374 100644
--- a/hw3-heat-fft/compute_temperature.cc
+++ b/hw3-heat-fft/compute_temperature.cc
@@ -1,85 +1,110 @@
 #include "compute_temperature.hh"
 #include "fft.hh"
 #include "material_point.hh"
 #include <cmath>
 
 /* -------------------------------------------------------------------------- */
 
 void ComputeTemperature::setDeltaT(Real dt) {
   this->dt = dt;
 }
 
 void ComputeTemperature::compute(System& system) {
 
     UInt size = system.getNbParticles();
     UInt N = std::sqrt(size);
 
     //////// VALUES TO BE SPECIFIED BY USER ////////////////////
-    Real dx = 0.1; // can also be extracted from distance between particle positions
+    Real L = 2.0;
+    Real dx = 0.2; // can also be extracted from distance between particle positions
     Real rho = 1.0;
     Real C = 1.0;
     Real kappa = 1.0;
     ///////////////////////////////////////////////////////////
 
     Real T = N*dx;
 
     // We have "size" particles aligned on a "sqrt(size) x sqrt(size)" grid
     Matrix<complex> theta(N);
     Matrix<complex> hv(N);
 
-    UInt k = 0;
+    //UInt k = 0;
     UInt i, j;
     for (auto& par : system) {
-          // Assume particles are already sorted in row-major format!
-          //        Make this a smarter loop
 
-          // Get 2d index from 1D index
-          j = N / k;
-          i = k % N;
+          MaterialPoint & m_par = dynamic_cast<MaterialPoint&>(par);
+
+          Vector pos = m_par.getPosition();
+          Real x = pos[0];
+          Real y = pos[1];
+
+          i = round(((x + L/2.0) / dx));
+          j = round(((y + L/2.0) / dx));
+
+          // Get 2d index from 1D index: k = i + j * N
+          //j = N / k;
+          //i = k % N;
 
           // Fill the matrix entry (corresponding to the particle par) with its temperature and heat rate
-          // theta(i,j) = par.getTemperature();
-          // hv(i,j)    = par.getHeatRate();
-          theta(i,j) = 1.0;
-          hv(i,j)    = 1.0;
+          theta(i,j) = m_par.getTemperature();
+          hv(i,j)    = m_par.getHeatRate();
+          //theta(i,j) = 1.0;
+          //hv(i,j)    = 1.0;
 
           // update 1D index
-          k++;
+          //k++;
     } // end for loop over particles
+
     std::cout << "We looped over this number of particles: " << std::endl;
 
     Matrix<complex> theta_hat = FFT::transform(theta);
 
     Matrix<std::complex<int>>    q_squared = FFT::computeFrequencies(N);
 
-    // Temporary solution by converting from complex-int to complex-double
+    /////// Temporary solution by converting from complex-int to complex-double
     Matrix<std::complex<double>> q_squared_converted(N);
     for (auto&& entry : index(q_squared_converted)) {
           int i = std::get<0>(entry);
           int j = std::get<1>(entry);
           std::complex<double> in = ( std::real(q_squared(i,j)), std::imag(q_squared(i,j)) );
           auto& out = std::get<2>(entry);
           out = in;
     }
+   //////////////////////////////////////////////////////////////////////////////
 
     Matrix<complex> d_theta_hat_dt(N);
     for (auto&& entry : index(d_theta_hat_dt)) {
           int i = std::get<0>(entry);
           int j = std::get<1>(entry);
           auto& val = std::get<2>(entry);
           val = ( 1 / (rho * C) ) * ( hv(i,j) - kappa * theta_hat(i,j) * (2*M_PI/T)*(2*M_PI/T) * q_squared_converted(i,j) );
     }
 
     Matrix<complex> d_theta_dt = FFT::itransform(d_theta_hat_dt);
 
     for (auto&& entry : index(theta)) {
           int i = std::get<0>(entry);
           int j = std::get<1>(entry);
           auto& val = std::get<2>(entry);
           val += dt * d_theta_dt(i,j);
     }
 
+    for (auto& par : system) {
+
+          MaterialPoint & m_par = dynamic_cast<MaterialPoint&>(par);
+
+          Vector pos = m_par.getPosition();
+          Real x = pos[0];
+          Real y = pos[1];
+
+          i = round(((x + L/2.0) / dx));
+          j = round(((y + L/2.0) / dx));
+
+          m_par.getTemperature() = std::real(theta(i,j));
+
+    } // end for loop over particles
+
 
 }
 
 /* -------------------------------------------------------------------------- */