Как писать научный и абстрактный код на Python

Многие пишут код на Python прямым способом без применения абстракций. Однако средства Python позволяют писать более обобщенный код (с учетом принципов DRY и SOLID), который может пригодиться в будущем. Сегодня пойдет речь о том, как это сделать на примере научных вычислений.

Стоящая задача

Наша задача продемонстрировать, как вычислять ожидаемое значение случайного процесса. Математически это выражается как

    \[ \mathbb{E}[X] = \sum_{i \in \mathcal{X}} X_i p(X_i) \]

где p(\mathcal{X}) — это функция вероятности, или для непрерывных значений:

    \[ \mathbb{E}[X] = \int_\mathcal{X} x p(x) dx \]

где p(x) — это плотность вероятности.

Возможно вам захочется работать с любым распределением: как дискретным, так и непрерывным. Как минимум, вы захотите узнать природу проблемы, чтобы под нее адаптировать свой код.

Стартовая точка — плохой код на Python

Начнем с плохого кода. Допустим бросаем шестигранный кубик. Поскольку падение на любую грань событие равновероятное, то здесь приходит в голову посчитать среднее значение за несколько бросков.

import random
import numpy as np


def die(sides=6):
    return random.randint(1, sides)

def expected_value(n_samples):
    samples = [die() for _ in range(n_samples)]
    return np.mean(samples)

Что не так с этим кодом? Во-первых, функция die возвращает одно испытание за раз. Поэтому она должна быть вызвана N раз, что уже замедляет процесс.

Во-вторых, expected_value функция сильно зависит от функции die, которая производит испытания. Это неудобно, ведь можно бросить и кубик с другим количеством граней. Тогда придется передавать в expected_value еще один параметр. Но интерфейс expected_value станет немного контринтуитивным, при этом она все еще опирается на конкретное явление, в данном случае кубик, с заданным распределение.

Разработка и внедрение ML-решений

Код курса
MLOPS
Ближайшая дата курса
1 июля, 2024
Продолжительность
24 ак.часов
Стоимость обучения
54 000 руб.

Улучшаем код в Python

Рассмотрим те решения, которые мы можем применить для этого случая.

Идея 1

Можно хранить дополнительную переменную для хранения испытаний. Пример кода на Python:

ef expected_value(samples):
    return np.mean(samples)

samples = [die(12) for _ in range(n_samples)]

ev = expected_value(samples)

Но глобальные переменные не самое лучшее решение. Их можно задеть, на них тратится память, они занимают идентификатор. А также ее нужно подготовить для функции expected_value.

Идея 2

Другое решение, при котором можно держать die внутри, — это передавать ее в виде объкта:

from functools import partial

twelve_sided_die = partial(die, 12)

def expected_values(die, n_samples):
    samples = [die() for _ in range(n_samples)]
    return np.mean(samples)


ev = expected_values(twelve_sided_die, 10000)

Идея заключается в подготовленной версии die, а expected_value будет использовать её в виде источника испытаний. Однако появляется новая проблема: функция expected_value совместима только с функцией die. Она не может вычислить другой распределение или, как минимум, посчитать среднее правильно.

Идея 3

Можно пойти по другому пути и выйти на более высокий абстрактный уровень и спроектировать хороший интерфейс.

На абстрактном уровне у нас есть:

  • некоторое распределение (это может быть кубик, монета, нормальное распределение или любое другое);
  • математическая операция для преобразования одних данных в другие (среднее, отклонение и т.д.);

Давайте сосредоточимся на абстрактном интерфейсе.

Распределение

Распределение выражается функцией (непрерывной, дискретной, бесконечной, конечной), из которой можно составить выборку. Мы можем использовать существующие формулы распределения Гаусса или Пуассона, но это может быть и какая-то своя, например, унаследованная от гистограммы.

Итак, можем заключить, что у нас есть некая абстракция для выражения распределения:

from abc import ABC, abstractmethod

class Distribution(ABC):

    @abstractmethod
    def sample(self):
        pass

Реализация

Поскольку метод sample у нас абстрактный, то нам нужно его реализовать в подклассах. Например, для нашего кубика мы можем написать такой код на Python:

import numpy as np


class Die(Distribution):
    def __init__(self, sides = 6):
        self.sides = sides

    def sample(self, n_samples):
        return np.random.randint(1, self.sides+1, size=n_samples)

Теперь мы можем создавать объекты, например, Die(12).sample(100) с 12-ю гранями и выборкой из 100. Мы также заменили списковые включения (list comprehension) на массивы numpy для повышения производительности.

Мы можем пойти дальше. Определить для него методы для более информативного вывода и для сравнения объектов:

    def __repr__(self):
        return f"{self.__class__.__name__}(sides={self.sides})"

    def __eq__(self, other):
        if isinstance(other, self.__class__):
            return self.sides == other.sides
        return False

Модуль Dataclass в Python

Реализовывать сразу четыре метода может быть излишним. Кроме того, текущая реализация может изменять объект через присваивания, например, die.sides = 20.

Поэтому мы можем воспользоваться модулем dataclass, о котором говорили тут, для уменьшения кода

from dataclasses import dataclass


@dataclass(frozen=True)
class Die(Distribution):
    sides: int = 6

    def sample(self, n):
        return np.random.randint(1, self.sides+1, size=n)

Добавив frozen=True, мы запретили присваивать атрибуты.

Теперь наша функция expected_value, должна принимать объект распределения в качестве параметра. Причем этот объект должен иметь метод sample.

Аннотации типов

Для более педантичного кода мы можем аннотировать возвращаемые типы. Действительно, результатом sample может быть последовательность любого типа.

Перепишем код на Python, добавив другие распределения и типы:

from typing import Generic, Sequence, TypeVar


D = TypeVar("D")

class Distribution(ABC, Generic[D]):

    @abstractmethod
    def sample(self, n: int) -> Sequence[D]:
        pass


@dataclass(frozen=True)
class Die(Distribution[int]):
    sides: int = 6

    def sample(self, n: int) -> Sequence[int]:
        return np.random.randint(1, self.sides+1, size=n)


@dataclass(frozen=True)
class Coin(Distribution[str]):
    outcomes: tuple = ("H", "T")
    fairness: float = 0.5

    def sample(self, n: int) -> Sequence[str]:
        p = (self.fairness, 1.0 - self.fairness)
        return np.random.choice(self.outcomes, size=n, p=p)


@dataclass(frozen=True)
class Gaussian(Distribution[float]):
    mu: float = 0.0
    sigma: float = 1.0

    def sample(self, n: int) -> Sequence[float]:
        np.random.normal(loc=self.mu, scale=self.sigma, size=n)

Через D определяется пользовательский тип. В конкретных классах мы указываем конкретный тип. Также абстрактный класс наследует Generic[D] для создания параметрического класса.

Тогда функция вычисления среднего должна выглядеть так:

# Работает только с float и int, но не со str
# Позже подкоректируем эту функцию
def expected_value(d: Distribution[float | int) -> float:
    return d.sample(n=n).mean()

Мы могли бы не быть настолько педантичными, а воспользоваться утиной типизацией, о которой говорили тут. Действительно, ведь функция требует лишь наличия метода sample. Тогда можно определить протокол:

from typing import Protocol, TypeVar


D = TypeVar("D")

class Distribution(Protocol, Generic[D]):
    @abstractmethod
    def sample(self, n: int) -> Sequence[D]:
        pass

Тогда будет достаточно написать:

def expected_value(d: Distribution, n: int) -> float:
    ...

Вычисления

Вернемся к задаче вычислений. У нас есть функция expected_value, которая работает для числовых распределений. Таким образом, мы можем посчитать мат.ожидание Die и Gaussian, но не Coin.

В этом случае у нас два пути:

  • создать приближенное распределение путем отображения, например, ("H", "T") -> (0, 1);
  • оформить отображение внутри функции с помощью вспомогательного “адаптера”;

Первый способ не очень гибкий, кто-то другой может переставить местами, например, ("H", "T") -> (1, 0), тогда получим неверные результаты.

Вместо этого преобразуем функцию с адаптером:

def expected_value(
    d: Distribution[D],
    f: Callable[[D], Any] = lambda x: x,
    n: int = 1000
) -> float:
    return np.mean(np.apply_along_axis(f, axis=0, arr=d.sample(n)))

Второй параметр — это функция, которой мы можем доверить отображение данных. По умолчанию она не будет как-то изменять данные.

Теперь можно писать так:

gaussian = Gaussian(mu=4.0, sigma=2.0)
expected_value(gaussian, n=100000)

# but
coin = Coin(fairness=0.75)
expected_value(coin, f=lambda x: np.where(x == "H", 1.0, 0.0))

Композиция и итерирование в Python

Абстракции могут быть полезными и для итеративных алгоритмов. Часто приходится брать в расчет промежуточные результаты, например: сходимость, ранний останов, значения метрик и т.д.

Например экспоненту можно посчитать через предел:

    \[ e = \lim_{n \to \infty} \left(1 + \frac{1}{n}\right)<sup>n</sup> \]

Чем больше n, тем точнее приближение.

Начнем с решения “в лоб”:

def approximate_e(
    initial_guess: float,
    max_iter: int = 10,
    epsilon: float = 0.01
) -> float:
    e_value = initial_guess
    for n in range(1, max_iter + 1):
        new_e_value = (1.0 + 1.0 / n) ** n
        if abs(new_e_value - e_value) < epsilon:
            return new_e_value
        e_value = new_e_value

    return new_e_value

Однако, во-первых, данная функция делает три вещи вместо одной. Сама формула выражается строкой под номером 8. Если мы захотим заменить объект приближения (например, на корень квадратный), то нам придется создавать ещё одну функцию, которая различается одной строчкой. А это нарушения принципа DRY (Don’t Repeat Yourself).

Также алгоритм выдает результат сразу. Вместо фокусировки на математике и получения значений, когда просят, он бросает данные к вам. Для сложых алгоритмов это приведет к проблеме памяти.

Добавляем абстракцию

Итак, перед нами стоит 3 вещи:

  • что-то, что мы ожидаем в качестве верного числа (непосредственно вычисление);
  • что-то для остановки, когда исчерпано заданное количество итераций (ранний останов);
  • что-то для остановки, если достигли нужного результата (сходимость);

Тогда мы можем разбить функцию на эти 3 части:

from typing import Iterator
import itertools


def approximate_e() -> Iterator[float]:
    n = 1
    while True:
        yield (1.0 + 1.0 / n) ** n
        n += 1

def early_stop(
    values: Iterator[float], max_iter: int = 100
) -> Iterator[float]:
    return itertools.islice(values, max_iter)

def convergence(
    values: Iterator[float], epsilon: float = 0.01
) -> Iterator[float]:
    for a, b in itertools.pairwise(values):
        yield a
        if (a - b) < epsilon:
            break

Данное решение реализует “ленивое” вычисление с помощью yield. Данные получаем ровно тогда, когда просим. Теперь не нужно заботиться о памяти.

Кроме того, эти функции могут быть протестированы отдельно друг от друга.

Отметим, что здесь мы использовали функцию pairwise (о ней тут), которая появилась только в Python 3.10. Её можно реализовать самостоятельно:

def pairwise(values: Iterator[D]) -> Iterator[Tuple[D, D]]:
    a = next(values, None)
    if a is None:
        return
    for b in values:
        yield a, b
        a = b

Теперь можем вычислить экспоненту очень просто:

values = approximate_e()
values = early_stop(values, max_iter=50)
values = convergence(values, epsilon=0.001)

for value in values:
    print("e becomes:", value)

Разработка и внедрение ML-решений

Код курса
MLOPS
Ближайшая дата курса
1 июля, 2024
Продолжительность
24 ак.часов
Стоимость обучения
54 000 руб.

О том, как писать абстракции в своем коде на Python вы узнаете на наших образовательных курсах в лицензированном учебном центре обучения и повышения квалификации руководителей и ИТ-специалистов (менеджеров, архитекторов, инженеров, администраторов, Data Scientist’ов и аналитиков Big Data) в Москве:

Источники
  1. Оригинал

Добавить комментарий

Поиск по сайту