Skip to content

brahmap.utilities.modify_numpy_context

Bases: object

A context manager that temporarily replaces np.linalg.norm with brahmap.math.parallel_norm.

Source code in brahmap/utilities/tools.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
class modify_numpy_context(object):
    """A context manager that temporarily replaces `np.linalg.norm` with
    `brahmap.math.parallel_norm`."""

    def __init__(self) -> None:
        self.parallel_norm = parallel_norm
        self.original_norm = np.linalg.norm

    def __enter__(self) -> None:
        setattr(np.linalg, "norm", self.parallel_norm)

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        setattr(np.linalg, "norm", self.original_norm)