Skip to content

mccube._utils

Helpful utilities used throughout MCCube.

nop ¤

nop(*args, **kwargs) -> None

Callable which accepts any arguments and does nothing.

Example 
result = nop(123, test="string")
# None
Source code in mccube/_utils.py 
 8
 9
10
11
12
13
14
15
16
17
def nop(*args, **kwargs) -> None:
    """Callable which accepts any arguments and does nothing.

    Example:
        ```python
        result = nop(123, test="string")
        # None
        ```
    """
    ...

pack_particles ¤

pack_particles(particles: Particles, weights: Weights | None) -> Particles | PackedParticles

Packs particles and normalized weights into a single Particles PyTree.

Example 
particles = {"y1": jnp.ones((10, 2))} # Shape: [n=10, d=2]
weights = {"y1": jnp.ones(10)} # Shape: [n=10]
packed_particles = mccube.pack_particles(particles, weights)
# {"y1": jnp.c_[particles["y1"], weights["y1"]i / 10]} Shape: [n=n, d=d+1]

weights = {"y1": None}
packed_particles = mccube.pack_particles(particles, weights)
# {"y1": particles["y1"]} Shape: [n=n, d=d]

Parameters:

Returns:

Source code in mccube/_utils.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def pack_particles(
    particles: Particles, weights: Weights | None
) -> Particles | PackedParticles:
    """Packs `particles` and normalized `weights` into a single
    [`Particles`][mccube._custom_types.Particles] PyTree.

    Example:
        ```python
        particles = {"y1": jnp.ones((10, 2))} # Shape: [n=10, d=2]
        weights = {"y1": jnp.ones(10)} # Shape: [n=10]
        packed_particles = mccube.pack_particles(particles, weights)
        # {"y1": jnp.c_[particles["y1"], weights["y1"]i / 10]} Shape: [n=n, d=d+1]

        weights = {"y1": None}
        packed_particles = mccube.pack_particles(particles, weights)
        # {"y1": particles["y1"]} Shape: [n=n, d=d]
        ```

    Args:
        particles: is a [`Particles`][mccube._custom_types.Particles] PyTree.
        weights: is either a [`Weights`][mccube._custom_types.Weights] PyTree or [`None`][].

    Returns:
        A [`Particles`][mccube._custom_types.Particles] PyTree where each leaf has
            dimension `d+1` and this new trailing dimension represents the [`Weights`][mccube._custom_types.Weights].
            If `weights` is [`None`][], then `particles` is returned unmodified.
    """
    return jtu.tree_map(
        lambda p, w: jnp.c_[p, jnp.squeeze(w) / jnp.sum(w)] if w is not None else p,
        particles,
        weights,
    )

unpack_particles ¤

unpack_particles(particles: Particles, weighted: bool = False) -> tuple[Particles, None] | tuple[UnpackedParticles, Weights]

Unpacks PackedParticles into separate PyTrees of Particles and Weights.

Example 
particles = jnp.ones((10,2)) # Shape: [n=10, d=2]
weights = jnp.full(10, 1/10) # Shape: [n=10]
packed_particles = pack_particles(particles, weights) # Shape: [n=n, d=d+1]
result = unpack_particles(packed_particles, weighted=True)
# (particles, weights) Shape: [n=n, d=d], [n=n]

result = unpack_particles(packed_particles)
# (jnp.c_[particles, weights], None) Shape: [n=n, d=d+1], None

Parameters:

Returns:

Source code in mccube/_utils.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def unpack_particles(
    particles: Particles, weighted: bool = False
) -> tuple[Particles, None] | tuple[UnpackedParticles, Weights]:
    """Unpacks [`PackedParticles`][mccube._custom_types.PackedParticles] into separate
    PyTrees of [`Particles`][mccube._custom_types.Particles] and [`Weights`][mccube._custom_types.Weights].

    Example:
        ```python
        particles = jnp.ones((10,2)) # Shape: [n=10, d=2]
        weights = jnp.full(10, 1/10) # Shape: [n=10]
        packed_particles = pack_particles(particles, weights) # Shape: [n=n, d=d+1]
        result = unpack_particles(packed_particles, weighted=True)
        # (particles, weights) Shape: [n=n, d=d], [n=n]

        result = unpack_particles(packed_particles)
        # (jnp.c_[particles, weights], None) Shape: [n=n, d=d+1], None
        ```

    Args:
        particles: is a [`Particles`][mccube._custom_types.Particles] PyTree.
        weighted: indicates if to treat the `particles` as [`PackedParticles`][mccube._custom_types.PackedParticles]
            (`weighted==True`), or standard [`Particles`][mccube._custom_types.Particles]
            (`weighted==False`).

    Returns:
        A tuple of the unpacked `particles` and weights (`weighted==True`), or a tuple
            of `particles` and `None` (`weighted==False`).
    """
    _particles = jtu.tree_map(lambda p: p[..., :-1] if weighted else p, particles)
    _weights = jtu.tree_map(lambda p: p[..., -1] if weighted else None, particles)
    return _particles, _weights

all_subclasses ¤

all_subclasses(cls: type) -> set[type]

Recursively identifies all subclasses of a class in the current scope.

Example 
result = all_subclasses(mccube.AbstractRegion)
# {mccube.GaussianRegion}
Source code in mccube/_utils.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def all_subclasses(cls: type) -> set[type]:
    """Recursively identifies all subclasses of a class in the current scope.

    Example:
        ```python
        result = all_subclasses(mccube.AbstractRegion)
        # {mccube.GaussianRegion}
        ```
    """
    subclasses = set()

    for subclass in cls.__subclasses__():
        subclasses.add(subclass)
        subclasses.update(all_subclasses(subclass))
    return subclasses