From 1c2fecf3a24e5a90918451fdd5a0ec3f013bff62 Mon Sep 17 00:00:00 2001
From: Rowan Cockett <rowan@3ptscience.com>
Date: Tue, 16 Feb 2016 22:07:33 -0800
Subject: [PATCH] Add the IRLS Directive.

---
 SimPEG/Directives.py | 50 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/SimPEG/Directives.py b/SimPEG/Directives.py
index 2ed27c20..476c5916 100644
--- a/SimPEG/Directives.py
+++ b/SimPEG/Directives.py
@@ -271,7 +271,6 @@ class SaveOutputDictEveryIteration(_SaveEveryIteration):
         np.savez('{:s}-{:03d}'.format(self.fileName,self.opt.iter), iter=self.opt.iter, beta=self.invProb.beta, phi_d=self.invProb.phi_d, phi_m=self.invProb.phi_m, phi_ms=phi_ms, phi_mx=phi_mx, phi_my=phi_my, phi_mz=phi_mz,f=self.opt.f, m=self.invProb.curModel,dpred=self.invProb.dpred)
 
 
-
 # class UpdateReferenceModel(Parameter):
 
 #     mref0 = None
@@ -283,3 +282,52 @@ class SaveOutputDictEveryIteration(_SaveEveryIteration):
 #             mref = self.mref0
 #         self.m_prev = self.invProb.m_current
 #         return mref
+
+class update_IRLS(InversionDirective):
+
+     m = None
+     eps_min = None
+     factor = None
+     gamma = None
+     phi_m_last = None
+
+     def initialize(self):
+
+         # Scale the regularization for changes in norm
+         if getattr(self, 'phi_m_last', None) is not None:
+             self.reg.gamma = 1.
+             phim_new = self.reg.eval(self.invProb.curModel)
+             self.gamma = self.phi_m_last / phim_new
+
+         self.reg.gamma = self.gamma
+
+     def endIter(self):
+         # Cool the threshold parameter
+         if getattr(self, 'factor', None) is not None:
+             eps = self.reg.eps / self.factor
+
+             if getattr(self, 'eps_min', None) is not None:
+                 self.reg.eps = np.max([self.eps_min,eps])
+             else:
+                 self.reg.eps = eps
+
+
+         # Update the model used for the IRLS weights
+         if getattr(self, 'm', None) is None:
+             self.reg.m = self.invProb.curModel
+
+         # Update the pre-conditioner
+         diagA = np.sum(self.prob.G**2.,axis=0) + self.invProb.beta*(self.reg.W.T*self.reg.W).diagonal() * (self.reg.mapping * np.ones(self.prob.mesh.nC))**2.
+         PC     = Utils.sdiag(diagA**-1.)
+
+         self.opt.approxHinv = PC
+
+         phim_new = self.reg.eval(self.invProb.curModel)
+         self.reg.gamma = self.reg.gamma * self.invProb.phi_m_last / phim_new
+
+#==============================================================================
+#          import pylab as plt
+#          plt.figure()
+#          ax = plt.subplot(221)
+#          self.prob.mesh.plotSlice(self.invProb.curModel, ax = ax, normal = 'Z', ind=-5, clim = (0, 0.005))
+#==============================================================================