An example of how to design Factorio factories using linear algebra

rocket_fuel.py

"""An example of how to design Factorio factories using linear algebra.

This program computes the number of oil refineries, chemical plants, and
assembly machine 2s needed to convert an abundant supply of crude oil and water
into 1.33 rocket fuel per second, the rate needed to satisfy one rocket silo in
the base game, not including modules.

Two abstract vector spaces are constructed: one representing products, and one
representing recipes. Each recipe is associated with a balanced equation over
products. By expressing this association as a matrix, we can express the
factory and its outputs as a heterogeneous system of linear equations and use
numpy's linear algebra solver to compute the number of required machines.

Accounting for inputs within this framework requires treating them as the
output of a recipe with no inputs. They could simply be omitted from the
relevant equations, but by modeling them as recipes that produce one unit per
second, the same math that counts machines works out to compute the rate at
which these products are consumed.

For the rocket fuel example, the output is:

    WATER 111
    CRUDE_OIL 161
    ADVANCED_OIL_PROCESSING 9
    HEAVY_OIL_CRACKING 3
    SOLID_FUEL_FROM_LIGHT_OIL 9
    SOLID_FUEL_FROM_PETROLEUM_GAS 5
    ROCKET_FUEL 27
"""
import abc
import dataclasses
import enum
import math
from typing import Optional, Self, SupportsFloat

import numpy


class VectorSpace(enum.IntEnum):
    """Abstract base class for a vector space with enumerated axes."""

    def unit(self) -> numpy.ndarray:
        """The unit vector associated with an enumerated axis."""
        dim = len(type(self))
        axis = self.value - 1
        return numpy.array([int(i == axis) for i in range(dim)])


class Product(VectorSpace):
    """An input, output, or intermediate product."""
    WATER = enum.auto()
    CRUDE_OIL = enum.auto()
    HEAVY_OIL = enum.auto()
    LIGHT_OIL = enum.auto()
    PETROLEUM_GAS = enum.auto()
    SOLID_FUEL = enum.auto()
    ROCKET_FUEL = enum.auto()


class Recipe(VectorSpace):
    """An identifier for crafting recipes."""
    WATER = enum.auto()
    CRUDE_OIL = enum.auto()
    ADVANCED_OIL_PROCESSING = enum.auto()
    HEAVY_OIL_CRACKING = enum.auto()
    SOLID_FUEL_FROM_LIGHT_OIL = enum.auto()
    SOLID_FUEL_FROM_PETROLEUM_GAS = enum.auto()
    ROCKET_FUEL = enum.auto()


@dataclasses.dataclass
class RecipeData:
    """The data associated with a crafting recipe, irrespective of machine."""
    inputs: dict[Product, SupportsFloat]
    outputs: dict[Product, SupportsFloat]
    base_time_s: SupportsFloat

    @classmethod
    def external_input(cls, product: Product) -> Self:
        return cls({}, {product: 1}, base_time_s=1)


RECIPE_DATA: dict[Recipe, RecipeData] = {
    Recipe.WATER: RecipeData.external_input(Product.WATER),
    Recipe.CRUDE_OIL: RecipeData.external_input(Product.CRUDE_OIL),
    Recipe.ADVANCED_OIL_PROCESSING: RecipeData(
        inputs={
            Product.CRUDE_OIL: 100,
            Product.WATER: 50,
        },
        outputs={
            Product.HEAVY_OIL: 25,
            Product.LIGHT_OIL: 45,
            Product.PETROLEUM_GAS: 55,
        },
        base_time_s=5,
    ),
    Recipe.HEAVY_OIL_CRACKING: RecipeData(
        inputs={
            Product.HEAVY_OIL: 40,
            Product.WATER: 30,
        },
        outputs={
            Product.LIGHT_OIL: 30,
        },
        base_time_s=2,
    ),
    Recipe.SOLID_FUEL_FROM_LIGHT_OIL: RecipeData(
        inputs={Product.LIGHT_OIL: 10},
        outputs={Product.SOLID_FUEL: 1},
        base_time_s=1,
    ),
    Recipe.SOLID_FUEL_FROM_PETROLEUM_GAS: RecipeData(
        inputs={Product.PETROLEUM_GAS: 20},
        outputs={Product.SOLID_FUEL: 1},
        base_time_s=1,
    ),
    Recipe.ROCKET_FUEL: RecipeData(
        inputs={
            Product.LIGHT_OIL: 10,
            Product.SOLID_FUEL: 10,
        },
        outputs={
            Product.ROCKET_FUEL: 1,
        },
        base_time_s=15,
    ),
}


class Machine(abc.ABC):
    """A crafting machine, including any speed and productivity bonuses."""

    @abc.abstractmethod
    def input_speed(self) -> float:
        ...

    @abc.abstractmethod
    def output_speed(self) -> float:
        ...


@dataclasses.dataclass
class BasicMachine(Machine):
    """A basic machine with no speed or productivity bonuses."""
    base_speed: float = 1.0

    def input_speed(self) -> float:
        return self.base_speed

    def output_speed(self) -> float:
        return self.base_speed


EXTERNAL_INPUT = BasicMachine()
OIL_REFINERY = BasicMachine()
CHEMICAL_PLANT = BasicMachine()
BLUE_ASSEMBLER = BasicMachine(base_speed=0.75)


def product_vector(products: dict[Product, SupportsFloat]) -> numpy.ndarray:
    """Construct a Product vector from slightly more convenient syntax."""
    return sum(p.unit() * count for p, count in products.items())


def recipe_equation(rx: Recipe, m: Machine) -> numpy.ndarray:
    """Express a recipe as a Product (per second) vector."""
    data = RECIPE_DATA[rx]
    inputs = m.input_speed() * product_vector(data.inputs)
    outputs = m.output_speed() * product_vector(data.outputs)
    return (outputs - inputs) / data.base_time_s


@dataclasses.dataclass
class System:
    """A system of crafting machines to solve for.

    Each recipe used in the system is associated with exactly one
    crafting machine. The machine determines the crafting speed. The
    outputs are specified as desired units per second.
    """
    machines: dict[Recipe, Machine]
    outputs: dict[Product, SupportsFloat]

    def solve(self) -> dict[Recipe, float]:
        """Compute the number of machines needed making each recipe.

        External inputs are reported in units per second.
        """
        matrix = numpy.column_stack(
            [recipe_equation(rx, m) for rx, m in self.machines.items()])
        coeffs = numpy.linalg.solve(matrix, product_vector(self.outputs))
        return dict(zip(self.machines.keys(), coeffs))


MY_SYSTEM = System(
    machines={
        Recipe.WATER: EXTERNAL_INPUT,
        Recipe.CRUDE_OIL: EXTERNAL_INPUT,
        Recipe.ADVANCED_OIL_PROCESSING: OIL_REFINERY,
        Recipe.HEAVY_OIL_CRACKING: CHEMICAL_PLANT,
        Recipe.SOLID_FUEL_FROM_LIGHT_OIL: CHEMICAL_PLANT,
        Recipe.SOLID_FUEL_FROM_PETROLEUM_GAS: CHEMICAL_PLANT,
        Recipe.ROCKET_FUEL: BLUE_ASSEMBLER,
    },
    outputs={
        Product.ROCKET_FUEL: 4/3,
    },
)


for rx, count in MY_SYSTEM.solve().items():
    print(rx.name, math.ceil(count))