simulate.py

# pip install feather-format msgpack pandas python-ron

import collections
import functools
import math
import re
import subprocess
import shlex
import time
import uuid

from pathlib import Path

import msgpack
import pandas
import pyron


TUPLE_FIELD_PATTERN = re.compile(r"_\d+:\s*")

old_namedtuple = collections.namedtuple
collections.namedtuple = functools.partial(old_namedtuple, rename=True)


def ron_from_str(s: str):
    return pyron.loads(s.replace("'", '"'), preserve_structs=True)


def ron_to_str(r) -> str:
    return "".join(
        TUPLE_FIELD_PATTERN.sub("", pyron.to_string(r)).split()
    ).replace(",)", ")").replace(",]", "]").replace("Create()", "Create").replace("Resume()", "Resume")


def config_with_pause(config, pause_before: float, pause_config: str, pause_lineages: str):
    pause = ron_from_str(f"""Pause(
        before: {pause_before},
        config: {pause_config!r},
        destiny: Bincode(file: {pause_lineages!r}),
        mode: Resume,
    )""")
    
    return config._replace(pause=pause)


def config_with_downscale(config, downscale_x: int, downscale_y: int):
    downscale = ron_from_str(f"""Downscale(
        x: {downscale_x},
        y: {downscale_y},
        non_self_dispersal: UnlikelyApproximation(likelihood: 0.05),
    )""")

    scenario = config.scenario._replace(downscale=downscale)

    return config._replace(scenario=scenario)


def config_with_event_skipping(config):
    algorithm = ron_from_str("EventSkipping()")

    return config._replace(algorithm=algorithm)


def config_with_reporter_resume(config):
    resume = ron_from_str("Resume")

    location_reporter = config.reporters[1].reporters[0]._replace(mode=resume)
    species_plugin = config.reporters[1]._replace(reporters=[location_reporter])
    reporters = config.reporters[:1] + [species_plugin]

    return config._replace(reporters=reporters)


def get_pause_times(config, percentiles: list) -> list:
    speciation = config.speciation
    turnover = 0.5

    return [
        -math.log(1-p)/(turnover*speciation) for p in percentiles
    ]


def downscale_lineage(lineage, downscale_x_before: int, downscale_y_before: int, downscale_x_after: int, downscale_y_after: int):
    (global_reference, last_event_time, (x, y, index)) = lineage

    unscale_x = x + (index % downscale_x_before)
    unscale_y = y + (index // downscale_x_before)

    downscale_x = unscale_x % downscale_x_after
    downscale_y = unscale_y % downscale_y_after

    downscale_index = downscale_x + (downscale_y * downscale_x_after)

    downscale_x = unscale_x - downscale_x
    downscale_y = unscale_y - downscale_y

    return (global_reference, last_event_time, (downscale_x, downscale_y, downscale_index))


def downscale_lineages(pause_lineages, downscale_x_before: int, downscale_y_before: int, downscale_x_after: int, downscale_y_after: int) -> None:
    with open(pause_lineages, "rb") as f:
        lineages = msgpack.load(f)
    
    lineages = [
        downscale_lineage(lineage, downscale_x_before, downscale_y_before, downscale_x_after, downscale_y_after)
        for lineage in lineages
    ]

    with open(pause_lineages, "wb") as f:
        msgpack.dump(lineages, f)


def run_simulation(config) -> None:
    subprocess.run(shlex.split(
        "cargo run --release --features almost-infinite-normal-dispersal-scenario,"
        + "almost-infinite-clark2dt-dispersal-scenario,almost-infinite-downscaled-scenario,"
        + f"gillespie-algorithms --quiet -- simulate '{ron_to_str(config)}'"
    ), check=True)

file_suffix = str(uuid.uuid4())

config = ron_from_str(f"""Simulate(
    speciation: {1e-10},
    sample: Sample(percentage: {0.415}),
    rng: Seed({42}),
    
    algorithm: Gillespie(),

    scenario: AlmostInfinite(
        sample: Rectangle(
            origin: Location(x: {1}, y: {1}),
            width: {10000},
            height: {10000},
        ),
        dispersal: Clark2Dt(
            shape_u: {5.0},
            tail_p: {1.0},
        ),
        downscale: None,
    ),

    pause: None,

    log: EventLog(directory: "I-solemnly-swear-that-I-am-up-to-no-good"),
    
    reporters: [
        Plugin(
            library: "target/release/deps/libnecsim_plugins_common.so",
            reporters: [Execution(), Progress()],
        ),
        Plugin(
            library: "target/release/deps/libnecsim_plugins_species.so",
            reporters: [
                LocationSpeciesFeather(
                    output: {("species." + file_suffix + ".feather")!r},
                    deduplication: None,
                    mode: Create,
                ),
            ],
        ),
    ],
)
""")

pause_times = get_pause_times(config, percentiles=[
    0.0000001, 0.000025, 0.001, 0.1, 0.999,
])

downscale_factor = 8

print(pause_times)

start_time = time.monotonic()

downscale = 1

for i, pause_time in enumerate(pause_times):
    pause_config = f"resume-{i}.{file_suffix}.ron"
    pause_lineages = f"lineages-{i}.{file_suffix}.mpk"

    config = config_with_pause(config, pause_before=pause_time, pause_config=pause_config, pause_lineages=pause_lineages)

    run_simulation(config)

    if i > 0:
        Path(f"lineages-{i-1}.{file_suffix}.mpk").unlink()

    if not Path(pause_config).exists():
        break

    with open(pause_config, "r") as f:
        config = ron_from_str(f"Simulate{f.read()}")

    Path(pause_config).unlink()
    
    config = config_with_event_skipping(config)
    config = config_with_reporter_resume(config)

    downscale *= downscale_factor

    downscale_lineages(
        pause_lineages,
        downscale_x_before=downscale//downscale_factor, downscale_y_before=downscale//downscale_factor,
        downscale_x_after=downscale, downscale_y_after=downscale,
    )

    config = config_with_downscale(config, downscale_x=downscale, downscale_y=downscale)

if i > 0:
    Path(f"lineages-{i}.{file_suffix}.mpk").unlink(missing_ok=True)

end_time = time.monotonic()

species = pandas.read_feather(f"species.{file_suffix}.feather")

print(f"There are {species['species'].nunique()} species from {species['count'].sum()} individuals")
print(f"Running the simulation took {round(end_time - start_time, 2)}s overall")

Path(f"species.{file_suffix}.feather").unlink()