3d coloured noise

noise.py

import numpy as np
from scipy.fftpack import fftn, ifftn


def get_noise(noise_type: str, shape: tuple[int, int, int]) -> np.ndarray:
    white_noise = np.random.randn(*shape)
    if noise_type == 'white':
        noise = white_noise
    else:
        f_transform = fftn(white_noise)
        kx, ky, kz = map(np.fft.fftfreq, shape)
        kx, ky, kz = np.meshgrid(kx, ky, kz, indexing='ij')
        k = np.sqrt(kx ** 2 + ky ** 2 + kz ** 2)
        k[0, 0, 0] = 1  # Avoid division by zero
        if noise_type == 'red':
            f_transform /= k ** 2
        elif noise_type == 'pink':
            f_transform /= k
        elif noise_type == 'green':
            f_transform = f_transform
        elif noise_type == 'blue':
            f_transform *= k
        elif noise_type == 'violet':
            f_transform *= k ** 0.5
        elif noise_type == 'gray':
            f_transform *= k ** 2
        else:
            raise ValueError(f"Unsupported noise type: {noise_type}")
        noise = np.real(ifftn(f_transform))
    noise = (noise - noise.min()) / (noise.max() - noise.min())
    return noise