Repository: simpeg/simpeg
Pull request #1349 (2024-02-04 17:22:02)

Irls refactor
#### Summary

Refactoring of the UpdateIRLS directive used to control the `Sparse` regularization. 

- Attribute name changes to snake_case format
- Remove unused attributes and change to instance attributes.
- Move responsibility for the scaling of spherical parameters into a new `SphericalDomain` directive.
- Move responsibility for the cooling schedule to a `BetaSchedule` directive.

#### PR Checklist
* [x ] If this is a work in progress PR, set as a Draft PR
* [x] Linted my code according to the [style guides](https://docs.simpeg.xyz/content/getting_started/contributing/code-style.html).
* [x] Added [tests](https://docs.simpeg.xyz/content/getting_started/practices.html#testing) to verify changes to the code.
* [x] Added necessary documentation to any new functions/classes following the
      expect [style](https://docs.simpeg.xyz/content/getting_started/practices.html#documentation).
* [x] Marked as ready for review (if this is was a draft PR), and converted 
      to a Pull Request
* [x ] Tagged ``@simpeg/simpeg-developers`` when ready for review.

#### Reference issue
Related to restructuring proposed on #1327 

#### What does this implement/fix?
Reduces complexity of current `Update_IRLS` class + PEP8 standard. 

#### Additional information
Clean up deprecation warnings on #1472 

domfournier: Sorry for the messy commit history. Pylint went a bit nuts on the imports.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
Attention: `245 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`716e221`)](https://app.codecov.io/gh/simpeg/simpeg/commit/716e221e83b986479037ae90ba9787f53d9b7334?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 82.30% compared to head [(`7d9603c`)](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 7.38%.

| [Files](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) | Patch % | Lines |
|---|---|---|
| [SimPEG/directives/\_regularization.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL2RpcmVjdGl2ZXMvX3JlZ3VsYXJpemF0aW9uLnB5) | 0.00% | [236 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [SimPEG/directives/\_\_init\_\_.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL2RpcmVjdGl2ZXMvX19pbml0X18ucHk=) | 20.00% | [4 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [SimPEG/directives/directives.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL2RpcmVjdGl2ZXMvZGlyZWN0aXZlcy5weQ==) | 25.00% | [3 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [SimPEG/inverse\_problem.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL2ludmVyc2VfcHJvYmxlbS5weQ==) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main   #1349       +/-   ##
==========================================
- Coverage   82.30%   7.38%   -74.92%     
==========================================
  Files         165     166        +1     
  Lines       25239   25287       +48     
==========================================
- Hits        20773    1868    -18905     
- Misses       4466   23419    +18953     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/simpeg/simpeg/pull/1349?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).

santisoler: > Sorry for the messy commit history. Pycharm went a bit nuts on the imports.

No worries! We can [Squash and Merge](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/incorporating-changes-from-a-pull-request/about-pull-request-merges#squash-and-merge-your-commits) and clean it up later.

I see a few tests failing. Do you want me to review it now @domfournier? Or are you still working on it?
domfournier: Hi @santisoler , yes let me fix those. I can ping you when I am done with it. Errors are mostly related to examples using deprecated arguments, which I should handle more gracefully right here.

Are we ok with updating the tests/examples for suppress the warnings after, on a separate PR?
santisoler: > Hi @santisoler , yes let me fix those. I can ping you when I am done with it. Errors are mostly related to examples using deprecated arguments, which I should handle more gracefully right here.

Sounds good!

> Are we ok with updating the tests/examples for suppress the warnings after, on a separate PR?

It would be ok to update the examples here as well. But we can totally handle them in a separate PR if you feel we need to change a lot of them.
domfournier: @santisoler I reverted the deletion of the init. You will have to explain to me why we don't need those anymore.
domfournier: > Thanks for addressing my previous comments, @domfournier. Just went through a second pass over this PR. I think it's getting in good shape. I left a few comments below.
> 
> My major considerations with this PR are:
> 
> * The `UpdateIRLS` contains a `BetaSchedule` directive, generating a nested structure of directives. I'd like it better if the `UpdateIRLS` implements its own beta scheduling as required instead of relying on a underlying directive object.

I would push back on this one. That is how it was before, and just ended up duplicating a lot of code. I kind of like the idea of re-using the mechanics of BetaSchedule without exposing it to the user. I don't really see the harm in that?

> * Missing tests. The new class is not being tested. We should add tests for it. We would also need to keep the old tests for the deprecated `Update_IRLS` class. We'll remove those when we delete that class though.

Yes can do.

> 
> BTW, are you planning to update the examples using `Update_IRLS` in a separate PR, right?

Next PR: #1472 
> 
> Let me know if you need any clarification on my comments or any help with this!


domfournier: @santisoler ready for round 3

domfournier: @santisoler What's the next step?
domfournier: Would be great to move this forward 
domfournier: Aight @santisoler , ready for another round.
santisoler: @domfournier, I just pushed a few changes to this branch. Most of them minor things or documentation-related stuff.

* I added the new classes to the API reference.
* I changed the default values for the `input_norms` in the dataclass to None, since we cannot be sure that we'll always have four objective functions. Moreover, the default value was a tuple, but the type hint was stating it as a list. The issue is that we cannot set mutable variables (lists for example) as default variables. I think having it set to None is the best compromise.
* I bumped the simpeg version in which we'll remove the `Update_IRLS` class to `v0.24.0`.


I have one more concern about the new class. Currently, it's a child of `BetaSchedule`, therefore it inherits the `coolingFactor` and `coolingRate` properties (although they are not being used) and the `endIter` is being overwritten. So I don't see the point of that inheritance. Moreover, I see that this inheritance could lead to issues. For example, if you define the directive as:

```python
irls_directive = directives.UpdateIRLS(coolingFactor=3)
```

you won't get any error, since the `coolingFactor` is taken to the `**kwargs` and the `irls_directive.coolingFactor` is set by the `BetaSchedule.__init__`. Nonetheless, the `irls_directive.cooling_factor` (the one that it'll actually be used by the directive) is set as the default value. So, the user might think that they are changing one of the parameters of the directive, while in fact the value of `3` they are using doesn't have any impact in the inversion.

If we make `UpdateIRLS` a child of  `InversionDirective` we would avoid that, because users will get an error after trying to pass `coolingFactor` to the constructor.

What do you think?

---

PS: One of the tests in `test_directives.py` is actually using `coolingFactor`. We might want to fix it:
https://github.com/simpeg/simpeg/blob/e9d86229c326a9e259c6d85b882c3090fdf93e58/tests/base/test_directives.py#L276-L286
domfournier: I think that you should just PR on this one with changes you want. It's getting pretty discouraging to continue working on this when the target keeps on moving.
domfournier: Iteration 7. Can we get a merge @lheagy @jcapriot @santisoler 
domfournier: @santisoler Not sure why the second set of codecov is stuck. Might have to force merge.

Patch:
diff --git a/simpeg/directives/__init__.py b/simpeg/directives/__init__.py
index 69081767d..bbffa1ea9 100644
--- a/simpeg/directives/__init__.py
+++ b/simpeg/directives/__init__.py
@@ -55,6 +55,8 @@ Directives related to sparse inversions
 .. autosummary::
    :toctree: generated/
 
+   UpdateIRLS
+   SphericalUnitsWeights
    Update_IRLS
 
 
@@ -108,7 +110,6 @@ from .directives import (
     SaveModelEveryIteration,
     SaveOutputEveryIteration,
     SaveOutputDictEveryIteration,
-    Update_IRLS,
     UpdatePreconditioner,
     Update_Wj,
     AlphasSmoothEstimate_ByEig,
@@ -116,6 +117,7 @@ from .directives import (
     ScalingMultipleDataMisfits_ByEig,
     JointScalingSchedule,
     UpdateSensitivityWeights,
+    Update_IRLS,
     ProjectSphericalBounds,
 )
 
@@ -125,6 +127,8 @@ from .pgi_directives import (
     PGI_AddMrefInSmooth,
 )
 
+from ._regularization import UpdateIRLS, SphericalUnitsWeights
+
 from .sim_directives import (
     SimilarityMeasureInversionDirective,
     SimilarityMeasureSaveOutputEveryIteration,
diff --git a/simpeg/directives/_regularization.py b/simpeg/directives/_regularization.py
new file mode 100644
index 000000000..686813ed3
--- /dev/null
+++ b/simpeg/directives/_regularization.py
@@ -0,0 +1,492 @@
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+from dataclasses import dataclass
+
+from ..maps import Projection
+from .directives import InversionDirective, UpdatePreconditioner, BetaSchedule
+from ..regularization import (
+    Sparse,
+    BaseSparse,
+    SmoothnessFirstOrder,
+    WeightedLeastSquares,
+)
+from ..utils import validate_integer, validate_float
+
+
+@dataclass
+class IRLSMetrics:
+    """
+    Data class to store metrics used by the IRLS algorithm.
+
+    Parameters
+    ----------
+    input_norms : list of floats or None
+        List of norms temporarily stored during the initialization.
+    irls_iteration_count : int
+        Number of IRLS iterations.
+    start_irls_iter : int or None
+        Iteration number when the IRLS process started.
+    f_old : float
+        Previous value of the regularization function.
+    """
+
+    input_norms: list[float] | None = None
+    irls_iteration_count: int = 0
+    start_irls_iter: int | None = None
+    f_old: float = 0.0
+
+
+class UpdateIRLS(InversionDirective):
+    """
+    Directive to control the IRLS iterations for :class:`~simpeg.regularization.Sparse`.
+
+    Parameters
+    ----------
+    cooling_rate: int
+        Number of iterations to cool beta.
+    cooling_factor: float
+        Factor to cool beta.
+    chifact_start: float
+        Starting chi factor for the IRLS iterations.
+    chifact_target: float
+        Target chi factor for the IRLS iterations.
+    irls_cooling_factor: float
+        Factor to cool the IRLS threshold epsilon.
+    f_min_change: float
+        Minimum change in the regularization function to continue the IRLS iterations.
+    max_irls_iterations: int
+        Maximum number of IRLS iterations.
+    misfit_tolerance: float
+        Tolerance for the target misfit.
+    percentile: float
+        Percentile of the function values used to determine the initial IRLS threshold.
+    verbose: bool
+        Print information to the screen.
+    """
+
+    def __init__(
+        self,
+        cooling_rate: int = 1,
+        cooling_factor: float = 2.0,
+        chifact_start: float = 1.0,
+        chifact_target: float = 1.0,
+        irls_cooling_factor: float = 1.2,
+        f_min_change: float = 1e-2,
+        max_irls_iterations: int = 20,
+        misfit_tolerance: float = 1e-1,
+        percentile: float = 100.0,
+        verbose: bool = True,
+        **kwargs,
+    ):
+        self._metrics: IRLSMetrics | None = None
+        self.cooling_rate = cooling_rate
+        self.cooling_factor = cooling_factor
+        self.chifact_start: float = chifact_start
+        self.chifact_target: float = chifact_target
+        self.irls_cooling_factor: float = irls_cooling_factor
+        self.f_min_change: float = f_min_change
+        self.max_irls_iterations: int = max_irls_iterations
+        self.misfit_tolerance: float = misfit_tolerance
+        self.percentile: float = percentile
+
+        super().__init__(
+            verbose=verbose,
+            **kwargs,
+        )
+
+    @property
+    def metrics(self) -> IRLSMetrics:
+        """Various metrics used by the IRLS algorithm."""
+        if self._metrics is None:
+            self._metrics = IRLSMetrics()
+        return self._metrics
+
+    @property
+    def max_irls_iterations(self) -> int:
+        """Maximum irls iterations."""
+        return self._max_irls_iterations
+
+    @max_irls_iterations.setter
+    def max_irls_iterations(self, value):
+        self._max_irls_iterations = validate_integer(
+            "max_irls_iterations", value, min_val=0
+        )
+
+    @property
+    def misfit_tolerance(self) -> float:
+        """Tolerance on deviation from the target chi factor, as a fractional percent."""
+        return self._misfit_tolerance
+
+    @misfit_tolerance.setter
+    def misfit_tolerance(self, value):
+        self._misfit_tolerance = validate_float("misfit_tolerance", value, min_val=0)
+
+    @property
+    def percentile(self) -> float:
+        """Tolerance on deviation from the target chi factor, as a fractional percent."""
+        return self._percentile
+
+    @percentile.setter
+    def percentile(self, value):
+        self._percentile = validate_float(
+            "percentile", value, min_val=0.0, max_val=100.0
+        )
+
+    @property
+    def chifact_start(self) -> float:
+        """Target chi factor to start the IRLS process."""
+        return self._chifact_start
+
+    @chifact_start.setter
+    def chifact_start(self, value):
+        self._chifact_start = validate_float(
+            "chifact_start", value, min_val=0, inclusive_min=False
+        )
+
+    @property
+    def chifact_target(self) -> float:
+        """Targer chi factor to maintain during the IRLS process."""
+        return self._chifact_target
+
+    @chifact_target.setter
+    def chifact_target(self, value):
+        self._chifact_target = validate_float(
+            "chifact_target", value, min_val=0, inclusive_min=False
+        )
+
+    @property
+    def cooling_factor(self):
+        """Beta is divided by this value every :attr:`cooling_rate` iterations.
+
+        Returns
+        -------
+        float
+        """
+        return self._cooling_factor
+
+    @cooling_factor.setter
+    def cooling_factor(self, value):
+        self._cooling_factor = validate_float(
+            "cooling_factor", value, min_val=0.0, inclusive_min=False
+        )
+
+    @property
+    def cooling_rate(self):
+        """Cool beta after this number of iterations.
+
+        Returns
+        -------
+        int
+        """
+        return self._cooling_rate
+
+    @cooling_rate.setter
+    def cooling_rate(self, value):
+        self._cooling_rate = validate_integer("cooling_rate", value, min_val=1)
+
+    @property
+    def irls_cooling_factor(self) -> float:
+        """IRLS threshold parameter (epsilon) is divided by this value every iteration."""
+        return self._irls_cooling_factor
+
+    @irls_cooling_factor.setter
+    def irls_cooling_factor(self, value):
+        self._irls_cooling_factor = validate_float(
+            "irls_cooling_factor", value, min_val=0.0, inclusive_min=False
+        )
+
+    @property
+    def f_min_change(self) -> float:
+        """Target chi factor to start the IRLS process."""
+        return self._f_min_change
+
+    @f_min_change.setter
+    def f_min_change(self, value):
+        self._f_min_change = validate_float(
+            "f_min_change", value, min_val=0, inclusive_min=False
+        )
+
+    def misfit_from_chi_factor(self, chi_factor: float) -> float:
+        """
+        Compute the target misfit from the chi factor.
+
+        Parameters
+        ----------
+        chi_factor : float
+            Chi factor to compute the target misfit from.
+        """
+        value = 0
+
+        for survey in self.survey:
+            value += survey.nD * chi_factor
+
+        return value
+
+    def adjust_cooling_schedule(self):
+        """
+        Adjust the cooling schedule based on the misfit.
+        """
+        ratio = self.invProb.phi_d / self.misfit_from_chi_factor(self.chifact_target)
+
+        if (
+            np.abs(1.0 - ratio) > self.misfit_tolerance
+            and self.metrics.start_irls_iter is not None
+        ):
+
+            if ratio > 1:
+                ratio = np.mean([2.0, ratio])
+            else:
+                ratio = np.mean([0.75, ratio])
+
+            self.cooling_factor = ratio
+
+    def initialize(self):
+        """
+        Initialize the IRLS iterations with l2-norm regularization (mode:1).
+        """
+
+        input_norms = []
+        for reg in self.reg.objfcts:
+            if not isinstance(reg, Sparse):
+                input_norms += [None]
+            else:
+                input_norms += [reg.norms]
+                reg.norms = [2.0 for _ in reg.objfcts]
+
+        self._metrics = IRLSMetrics(input_norms=input_norms)
+
+    def endIter(self):
+        """
+        Check on progress of the inversion and start/update the IRLS process.
+        """
+        # After reaching target misfit with l2-norm, switch to IRLS (mode:2)
+        if (
+            self.metrics.start_irls_iter is None
+            and self.invProb.phi_d < self.misfit_from_chi_factor(self.chifact_start)
+        ):
+            self.start_irls()
+
+        # Check if misfit is within the tolerance, otherwise scale beta
+        self.adjust_cooling_schedule()
+
+        # Only update after GN iterations
+        if (
+            self.metrics.start_irls_iter is not None
+            and (self.opt.iter - self.metrics.start_irls_iter) % self.cooling_rate == 0
+        ):
+            if self.stopping_criteria():
+                self.opt.stopNextIteration = True
+                return
+            else:
+                self.opt.stopNextIteration = False
+
+            # Print to screen
+            for reg in self.reg.objfcts:
+                if not isinstance(reg, Sparse):
+                    continue
+
+                for obj in reg.objfcts:
+                    if isinstance(reg, (Sparse, BaseSparse)):
+                        obj.irls_threshold /= self.irls_cooling_factor
+
+            self.metrics.irls_iteration_count += 1
+
+            # Reset the regularization matrices so that it is
+            # recalculated for current model. Do it to all levels of comboObj
+            for reg in self.reg.objfcts:
+                if not isinstance(reg, Sparse):
+                    continue
+
+                reg.update_weights(reg.model)
+
+            self.invProb.phi_m_last = self.reg(self.invProb.model)
+
+        # Apply beta cooling schedule mechanism
+        if self.opt.iter > 0 and self.opt.iter % self.cooling_rate == 0:
+            self.invProb.beta /= self.cooling_factor
+
+    def start_irls(self):
+        """
+        Start the IRLS iterations by computing the initial threshold values.
+        """
+        if self.verbose:
+            print(
+                "Reached starting chifact with l2-norm regularization:"
+                + " Start IRLS steps..."
+            )
+
+        self.metrics.start_irls_iter = getattr(self.opt, "iter", 0)
+        self.invProb.phi_m_last = self.reg(self.invProb.model)
+
+        # Either use the supplied irls_threshold, or fix base on distribution of
+        # model values
+        for reg, norms in zip(self.reg.objfcts, self.metrics.input_norms):
+            if not isinstance(reg, Sparse):
+                continue
+
+            for obj in reg.objfcts:
+                threshold = np.percentile(
+                    np.abs(obj.mapping * obj._delta_m(self.invProb.model)),
+                    self.percentile,
+                )
+                if isinstance(obj, SmoothnessFirstOrder):
+                    threshold /= reg.regularization_mesh.base_length
+
+                obj.irls_threshold = threshold
+
+            reg.norms = norms
+
+            if self.verbose:
+                print("irls_threshold " + str(reg.objfcts[0].irls_threshold))
+
+        # Save l2-model
+        self.invProb.l2model = self.invProb.model.copy()
+
+    def validate(self, directiveList=None):
+        directive_list = directiveList.dList
+        self_ind = directive_list.index(self)
+        lin_precond_ind = [isinstance(d, UpdatePreconditioner) for d in directive_list]
+
+        if any(lin_precond_ind) and lin_precond_ind.index(True) < self_ind:
+            raise AssertionError(
+                "The directive 'UpdatePreconditioner' must be after Update_IRLS "
+                "in the directiveList"
+            )
+        else:
+            warnings.warn(
+                "Without a Linear preconditioner, convergence may be slow. "
+                "Consider adding `directives.UpdatePreconditioner` to your "
+                "directives list",
+                stacklevel=2,
+            )
+
+        beta_schedule = [
+            d for d in directive_list if isinstance(d, BetaSchedule) and d is not self
+        ]
+
+        if beta_schedule:
+            raise AssertionError(
+                "Beta scheduling is handled by the `UpdateIRLS` directive."
+                "Remove the redundant `BetaSchedule` from your list of directives.",
+            )
+
+        spherical_scale = [isinstance(d, SphericalUnitsWeights) for d in directive_list]
+        if any(spherical_scale):
+            assert spherical_scale.index(True) < self_ind, (
+                "The directive 'SphericalUnitsWeights' must be before UpdateIRLS "
+                "in the directiveList"
+            )
+
+        return True
+
+    def stopping_criteria(self):
+        """
+        Check for stopping criteria of max_irls_iteration or minimum change.
+        """
+        phim_new = 0
+        for reg in self.reg.objfcts:
+            if isinstance(reg, (Sparse, BaseSparse)):
+                reg.model = self.invProb.model
+                phim_new += reg(reg.model)
+
+        # Check for maximum number of IRLS cycles
+        if self.metrics.irls_iteration_count == self.max_irls_iterations:
+            if self.verbose:
+                print(
+                    "Reach maximum number of IRLS cycles:"
+                    + f" {self.max_irls_iterations:d}"
+                )
+            return True
+
+        # Check if the function has changed enough
+        f_change = np.abs(self.metrics.f_old - phim_new) / (self.metrics.f_old + 1e-12)
+
+        if (
+            f_change < self.f_min_change
+            and self.metrics.irls_iteration_count > 1
+            and np.abs(
+                1.0
+                - self.invProb.phi_d / self.misfit_from_chi_factor(self.chifact_target)
+            )
+            < self.misfit_tolerance
+        ):
+            if self.verbose:
+                print("Minimum decrease in regularization. End of IRLS")
+            return True
+
+        self.metrics.f_old = phim_new
+
+        return False
+
+
+class SphericalUnitsWeights(InversionDirective):
+    """
+    Directive to update the regularization weights to account for spherical
+    parameters in radian and SI.
+
+    The scaling applied to the regularization weights is based on the ratio
+    between the maximum value of the model and the maximum value of angles (pi).
+
+    Parameters
+    ----------
+    amplitude: Projection
+        Map to the model parameters for the amplitude of the vector
+    angles: list[WeightedLeastSquares]
+        List of WeightedLeastSquares for the angles.
+    verbose: bool
+        Print information to the screen.
+    """
+
+    def __init__(
+        self,
+        amplitude: Projection,
+        angles: list[WeightedLeastSquares],
+        verbose: bool = True,
+        **kwargs,
+    ):
+
+        if not isinstance(amplitude, Projection):
+            raise TypeError(
+                "Attribute 'amplitude' must be of type " "'wires.Projection'"
+            )
+
+        self._amplitude = amplitude
+
+        if not isinstance(angles, (list, tuple)) or not all(
+            [isinstance(fun, WeightedLeastSquares) for fun in angles]
+        ):
+            raise TypeError(
+                "Attribute 'angles' must be a list of "
+                "'regularization.WeightedLeastSquares'."
+            )
+
+        self._angles = angles
+
+        super().__init__(
+            verbose=verbose,
+            **kwargs,
+        )
+
+    def initialize(self):
+        self.update_scaling()
+
+    def endIter(self):
+        self.update_scaling()
+
+    def update_scaling(self):
+        """
+        Add an 'angle_scale' to the list of weights on the angle regularization for the
+        different block of models to account for units of radian and SI.
+        """
+        amplitude = self._amplitude * self.invProb.model
+        max_p = max(amplitude)
+
+        for reg in self._angles:
+            for obj in reg.objfcts:
+                if obj.units != "radian":
+                    continue
+
+                obj.set_weights(angle_scale=np.ones_like(amplitude) * max_p / np.pi)
diff --git a/simpeg/directives/directives.py b/simpeg/directives/directives.py
index 75ad4834a..22fd36dc4 100644
--- a/simpeg/directives/directives.py
+++ b/simpeg/directives/directives.py
@@ -1,4 +1,7 @@
 from __future__ import annotations  # needed to use type operands in Python 3.8
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 import matplotlib.pyplot as plt
 import warnings
@@ -6,7 +9,7 @@ import os
 import scipy.sparse as sp
 from ..typing import RandomSeed
 from ..data_misfit import BaseDataMisfit
-from ..objective_function import ComboObjectiveFunction
+from ..objective_function import BaseObjectiveFunction, ComboObjectiveFunction
 from ..maps import IdentityMap, Wires
 from ..regularization import (
     WeightedLeastSquares,
@@ -32,6 +35,7 @@ from ..utils import (
     validate_string,
 )
 from ..utils.code_utils import (
+    deprecate_class,
     deprecate_property,
     validate_type,
     validate_integer,
@@ -39,6 +43,10 @@ from ..utils.code_utils import (
     validate_ndarray_with_shape,
 )
 
+if TYPE_CHECKING:
+    from ..simulation import BaseSimulation
+    from ..survey import BaseSurvey
+
 
 class InversionDirective:
     """Base inversion directive class.
@@ -140,7 +148,7 @@ class InversionDirective:
         return self.invProb.opt
 
     @property
-    def reg(self):
+    def reg(self) -> BaseObjectiveFunction:
         """Regularization associated with the directive.
 
         Returns
@@ -164,7 +172,7 @@ class InversionDirective:
         self._reg = value
 
     @property
-    def dmisfit(self):
+    def dmisfit(self) -> BaseObjectiveFunction:
         """Data misfit associated with the directive.
 
         Returns
@@ -188,7 +196,7 @@ class InversionDirective:
         self._dmisfit = value
 
     @property
-    def survey(self):
+    def survey(self) -> list[BaseSurvey]:
         """Return survey for all data misfits
 
         Assuming that ``dmisfit`` is always a ``ComboObjectiveFunction``,
@@ -203,7 +211,7 @@ class InversionDirective:
         return [objfcts.simulation.survey for objfcts in self.dmisfit.objfcts]
 
     @property
-    def simulation(self):
+    def simulation(self) -> list[BaseSimulation]:
         """Return simulation for all data misfits.
 
         Assuming that ``dmisfit`` is always a ``ComboObjectiveFunction``,
@@ -1963,6 +1971,7 @@ class SaveOutputDictEveryIteration(SaveEveryIteration):
         self.outDict[self.opt.iter] = iterDict
 
 
+@deprecate_class(removal_version="0.24.0", error=False)
 class Update_IRLS(InversionDirective):
     f_old = 0
     f_min_change = 1e-2
diff --git a/simpeg/inverse_problem.py b/simpeg/inverse_problem.py
index e554c95ce..8bc52a24c 100644
--- a/simpeg/inverse_problem.py
+++ b/simpeg/inverse_problem.py
@@ -96,7 +96,7 @@ class BaseInvProblem:
         self._counter = value
 
     @property
-    def dmisfit(self):
+    def dmisfit(self) -> ComboObjectiveFunction:
         """The data misfit.
 
         Returns
@@ -113,7 +113,7 @@ class BaseInvProblem:
         self._dmisfit = value
 
     @property
-    def reg(self):
+    def reg(self) -> ComboObjectiveFunction:
         """The regularization object for the inversion
 
         Returns
diff --git a/tutorials/03-gravity/plot_inv_1b_gravity_anomaly_irls.py b/tutorials/03-gravity/plot_inv_1b_gravity_anomaly_irls.py
index 7bfd6da31..5bb75bfa9 100644
--- a/tutorials/03-gravity/plot_inv_1b_gravity_anomaly_irls.py
+++ b/tutorials/03-gravity/plot_inv_1b_gravity_anomaly_irls.py
@@ -271,13 +271,12 @@ starting_beta = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
 
 # Defines the directives for the IRLS regularization. This includes setting
 # the cooling schedule for the trade-off parameter.
-update_IRLS = directives.Update_IRLS(
+update_IRLS = directives.UpdateIRLS(
     f_min_change=1e-4,
     max_irls_iterations=30,
-    coolEpsFact=1.5,
-    beta_tol=1e-2,
+    irls_cooling_factor=1.5,
+    misfit_tolerance=1e-2,
 )
-
 # Options for outputting recovered models and predicted data for each beta.
 save_iteration = directives.SaveOutputEveryIteration(save_txt=False)
 


Test patch:
diff --git a/tests/base/test_directives.py b/tests/base/test_directives.py
index f6450eb58..a47ae5720 100644
--- a/tests/base/test_directives.py
+++ b/tests/base/test_directives.py
@@ -116,14 +116,14 @@ class ValidationInInversion(unittest.TestCase):
         with self.assertRaises(AssertionError):
             # validation should happen and this will fail
             # (IRLS needs to be before update_Jacobi)
-            inv = inversion.BaseInversion(
+            inversion.BaseInversion(
                 invProb, directiveList=[betaest, update_Jacobi, IRLS]
             )
 
         with self.assertRaises(AssertionError):
             # validation should happen and this will fail
             # (sensitivity_weights needs to be before betaest)
-            inv = inversion.BaseInversion(
+            inversion.BaseInversion(
                 invProb, directiveList=[betaest, sensitivity_weights]
             )
 
@@ -258,7 +258,118 @@ class ValidationInInversion(unittest.TestCase):
 
         # self.test_sensitivity_weighting_subroutine(test_weights, test_directive)
 
-        print("SENSITIVITY WEIGHTING BY AMPLIUTDE AND MAX ALUE TEST PASSED")
+        print("SENSITIVITY WEIGHTING BY AMPLITUDE AND MAX VALUE TEST PASSED")
+
+    def test_irls_directive(self):
+        input_norms = [0.0, 1.0, 1.0, 1.0]
+        reg = regularization.Sparse(self.mesh)
+        reg.norms = input_norms
+        projection = maps.Projection(self.mesh.n_cells, np.arange(self.mesh.n_cells))
+
+        other_reg = regularization.WeightedLeastSquares(self.mesh)
+
+        invProb = inverse_problem.BaseInvProblem(self.dmis, reg + other_reg, self.opt)
+
+        beta_schedule = directives.BetaSchedule(coolingFactor=3)
+
+        # Here is where the norms are applied
+        irls_directive = directives.UpdateIRLS(
+            cooling_factor=3,
+            chifact_start=100.0,
+            chifact_target=1.0,
+            irls_cooling_factor=1.2,
+            f_min_change=np.inf,
+            max_irls_iterations=20,
+            misfit_tolerance=1e-0,
+            percentile=100,
+            verbose=True,
+        )
+
+        assert irls_directive.cooling_factor == 3
+        assert irls_directive.metrics is not None
+
+        # TODO Move these assertion test to the 'test_validation_in_inversion' after update
+        with self.assertRaises(AssertionError):
+            inversion.BaseInversion(
+                invProb, directiveList=[beta_schedule, irls_directive]
+            )
+
+        with self.assertRaises(AssertionError):
+            inversion.BaseInversion(
+                invProb, directiveList=[beta_schedule, irls_directive]
+            )
+
+        spherical_weights = directives.SphericalUnitsWeights(projection, [reg])
+        with self.assertRaises(AssertionError):
+            inversion.BaseInversion(
+                invProb, directiveList=[irls_directive, spherical_weights]
+            )
+
+        update_Jacobi = directives.UpdatePreconditioner()
+        with self.assertRaises(AssertionError):
+            inversion.BaseInversion(
+                invProb, directiveList=[update_Jacobi, irls_directive]
+            )
+
+        invProb.phi_d = 1.0
+        self.opt.iter = 3
+        invProb.model = np.random.randn(reg.regularization_mesh.nC)
+        inv = inversion.BaseInversion(invProb, directiveList=[irls_directive])
+
+        irls_directive.initialize()
+        assert irls_directive.metrics.input_norms == [input_norms, None]
+        assert reg.norms == [2.0, 2.0, 2.0, 2.0]
+
+        irls_directive.inversion = inv
+        irls_directive.endIter()
+
+        assert irls_directive.metrics.start_irls_iter == self.opt.iter
+        assert len(reg.objfcts[0]._weights) == 2  # With irls weights
+        assert len(other_reg.objfcts[0]._weights) == 1  # No irls
+        irls_directive.metrics.irls_iteration_count += 1
+        irls_directive.endIter()
+
+        assert self.opt.stopNextIteration
+
+        # Test stopping criteria based on max_irls_iter
+        irls_directive.max_irls_iterations = 2
+        assert irls_directive.stopping_criteria()
+
+        # Test beta re-adjustment down
+        invProb.phi_d = 4.0
+        irls_directive.misfit_tolerance = 0.1
+        irls_directive.adjust_cooling_schedule()
+        assert irls_directive.cooling_factor == 2.0
+
+        # Test beta re-adjustment up
+        invProb.phi_d = 0.5
+        irls_directive.adjust_cooling_schedule()
+        assert irls_directive.cooling_factor == 0.5
+
+    def test_spherical_weights(self):
+        reg = regularization.Sparse(self.mesh)
+        projection = maps.Projection(self.mesh.n_cells, np.arange(self.mesh.n_cells))
+        for obj in reg.objfcts[1:]:
+            obj.units = "radian"
+
+        with pytest.raises(TypeError, match="Attribute 'amplitude' must be of type"):
+            directives.SphericalUnitsWeights(reg, [reg])
+
+        with pytest.raises(TypeError, match="Attribute 'angles' must be a list of"):
+            directives.SphericalUnitsWeights(projection, ["abc"])
+
+        spherical_weights = directives.SphericalUnitsWeights(projection, [reg])
+
+        inv_prob = inverse_problem.BaseInvProblem(self.dmis, reg, self.opt)
+        model = np.abs(np.random.randn(reg.regularization_mesh.nC))
+        inv_prob.model = model
+        inv = inversion.BaseInversion(inv_prob, directiveList=[spherical_weights])
+        spherical_weights.inversion = inv
+
+        spherical_weights.initialize()
+
+        assert "angle_scale" not in reg.objfcts[0]._weights
+        assert reg.objfcts[1]._weights["angle_scale"].max() == model.max() / np.pi
 
     def tearDown(self):
         # Clean up the working directory