Computes the Geothmetic Meandian https://xkcd.com/2435/

geothmetic_meandian.py

#!/usr/bin/env python3
"""Module doc"""

import numpy as np
from sys import argv


def residual_size(values, axis=None):
    # shape: (4, 3, 5, 2, 3)
    # axis: (0, 2)
    # residual: 3*2*3 = 18
    # After selecting a sample on axis 0 and 2, it only has 12 dimensions left

    if axis is not None:
        return np.prod([s for i, s in enumerate(values.shape) if i not in axis])
    else:
        return values.size


def geo_mean(values, axis=None):
    arr = np.log(values)

    return np.exp(arr.sum(axis=axis) / residual_size(arr, axis=axis))


def F(values: np.ndarray, axis=None):
    arithm = values.mean(axis=axis)
    geom = geo_mean(values, axis=axis)
    median = np.median(values, axis=axis)

    assert arithm.shape == geom.shape == median.shape

    return np.array([arithm, geom, median])


def main():
    """Main function"""
    arr = np.random.rand(min(int(argv[1]), int(2e8)))
    arr = np.array([1, 1, 2, 3, 5])
    axis = None

    print("     arithm   | geometric | median    ")
    print("---------------------------------------")
    acc = F(arr, axis)
    for i in range(30000):
        print(f"{i+1:02d}# {acc}")
        new_acc = F(acc, axis=None)
        if np.allclose(new_acc - acc, 0, rtol=1e-9, atol=1e-9):
            print(f"{i+1:02d}# {new_acc}")
            print(f"Converges after {i+1} iterations")
            break
        acc = new_acc
    print(acc)

if __name__ == "__main__":
    main()