<jupyter_start><jupyter_text>Using Dask to parallelize simulationsIn real-world use, `SWYFT` will ofen be combined with existing simulators. As inference using truncated marginal neural ratio estimation (as implemented in `SWYFT`), allows the reuse **and** independent realization of simulations, SWYFT makes the use of more computationally expensive and (physically) relevant simulators viable. In `SWYFT` this supported by the use of the directory store, but even more importantly by the `DaskSimulator` class. In contrast to the `Simulator` class, `DaskSimulator` allows the simulations requested to executed in parallel on a Dask cluster.In this notebook we demonstrate the framework for the use of a more computationally expensive external simulator, although we still make use of the toy external simulator we previously defined.Here we have chosen to include all explanatory text.<jupyter_code>import numpy as np
import pylab as plt
import swyft
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Ntrain =3000
Npars = 2<jupyter_output><empty_output><jupyter_text>Set input ... In order to make use of en external simulator called frrom the command line, the user must specify a function to setup the simulator input. It should take one input argument (the array with the input parameters), and return any input to be passed to the program via `stdin`. If the simulator requires any input files to be present, this function should write these to disk.<jupyter_code>def set_input(v):
    v0 = v[0]
    v1 = v[1]
    v_str = str(v0).strip()+' '+str(v1).strip()
    return v_str<jupyter_output><empty_output><jupyter_text>... output methods ...Analogously, the user must define a function to retrieve results from the simulator output. It should take two input arguments (stdout and stderr of the simulator run) and return a dictionary with the simulator output shaped as described by the ``sim_shapes`` argument. If the simulator writes output to disk, this function **should** parse the results from the file(s).<jupyter_code>def get_output(stdout,stderr):
    try:
        if not stderr :      
            x0,x1 = stdout.split(" ")
            x0 = np.float(x0.strip())
            x1 = np.float(x1.strip())
            x = np.array([x0,x1])
            return dict(x=x)

        else:
            raise('simulator returned on stderr')

    except:
        raise('Error in output retrieval')

import os
cwd = os.getcwd()<jupyter_output><empty_output><jupyter_text>... and invocationHere we use the cell magic `%%writefile` command to create an external python function randgauss.py containing the simulator defined as model in the Quickstart notebook. This function is then invoked from the command line.<jupyter_code>%%writefile randgauss.py 
#!/usr/bin/env python

import numpy as np
import sys


def rgmodel(v,sigma=0.05):
    x = v + np.random.randn(2)*sigma
    return x


def main():
    sigma = None
    args = sys.stdin.readline()
    arg1, arg2 = args.split(' ')
    try:
        v0 = float(arg1.rstrip())
        v1 = float(arg2.rstrip())

    except:
        raise()

    v = np.array([v0,v1])

    if sigma is not None:
        x = rgmodel(v,sigma=sigma)
    else:
        x = rgmodel(v)

    print(str(x[0]).strip()+' '+str(x[1]).strip())


if __name__ == "__main__":
    main()<jupyter_output>Writing randgauss.py<jupyter_text>It is up to the user to ensure adaquate permissions for all relevant files.<jupyter_code>!chmod 755 randgauss.py<jupyter_output><empty_output><jupyter_text>And to ensure that the root temporary directory in which the simulator is run exists.   Each instance of the simulator will run in a separate sub-folder.<jupyter_code>!mkdir -p ./tmp
command = cwd+'/randgauss.py'<jupyter_output><empty_output><jupyter_text>Defining the Dask simulatorThe simulator itslef can then be defined using the `from_command()` method of the `DaskSimulator` class, exactly as for the `Simulator` class. A local dask cluster with a local scheduler is instantiated.<jupyter_code>simulator = swyft.DaskSimulator.from_command(
    command=command,
    pnames=["x0","x1"],
    sim_shapes=dict(x=(Npars,)),
    set_input_method=set_input,
    get_output_method=get_output,
    tmpdir=cwd+'/tmp/'
)<jupyter_output><empty_output><jupyter_text>In order to connect to an external Dask cluster, and also locally for better performance, we make use of the Dask distributed scheduler.<jupyter_code>from dask.distributed import LocalCluster # from dask_jobqueue import SLURMCluster<jupyter_output><empty_output><jupyter_text>Here we set the cluster the simulator will use, enabling the user to connect to a dask cluster of their choice.Please NoteWhen using the DaskSimulator class in combination with a simulator invoked from the command line, the user must ensure that it is available at the specified location on all nodes of the cluster.Furthermore, the number of threads per worker **must** be set to one, i.e. `threads_per_worker=1`. This is because each worker access a different scratch directory space, which can be achieved by process-parallelism only. The command-line simulator can then spawn multiple threads.<jupyter_code>cluster = LocalCluster(n_workers=3,threads_per_worker=1)
simulator.set_dask_cluster(cluster)<jupyter_output><empty_output><jupyter_text>The directory storeA directory store is initialized using the `Store.directory_store()` convenience function. The user must specify the path where to create a new store. If a store already exists, it can be opened using `Store.load()`.<jupyter_code>store = swyft.Store.directory_store(cwd+'/mystore',simulator=simulator)<jupyter_output>Creating new store.<jupyter_text>InferrenceWe can now define the intial prior, and given an observation, begin with the inferrence problem,<jupyter_code>v0 = np.zeros(Npars)
obs0 = {"x":[0.,0.]}
prior = swyft.Prior(lambda u: u*2-1, Npars)
store.add(Ntrain, prior)<jupyter_output>Store: Adding 478 new samples to simulator store.<jupyter_text>invoking the simulator to produce the required simulations. Here, we wait for the simmulations to complete (blocking the flow of execution) and then continue<jupyter_code>store.simulate()
bound = None
dataset = swyft.Dataset(Ntrain, prior, store, bound=bound)
    
post = swyft.Posteriors(dataset)
post.add([(0, 1)], device=DEVICE)
post.train([(0, 1)], max_epochs = 20)
    
new_bound = post.truncate([(0, 1)], obs0)<jupyter_output>Bounds: Truncating...
Bounds: ...done. New volue is V=0.1877<jupyter_text>We can now use the truncated pior given by `new_bound` to zoom in on the region(s) of interest for the inferrence by repeating the steps above.<jupyter_code>store.add(Ntrain,prior,bound = new_bound)<jupyter_output>Store: Adding 403 new samples to simulator store.<jupyter_text>Asynchronous simulationIn some cases performing all the required simulations, even when parallelized, will take a significant amount of time. For such situations, the use of the `DaskSimulator` class, together with the directory store enables the asynchronous execution of the requested simulations. The dask workers then write the obtained simulation results directly to the store on disk, allowing the inference workflow to be suspened and later restarted. **Please note**For the example shown here, we have started the dask cluster from within the notebook. If the user expects simulations to take so much time as to warrant exiting and rejoining at a later time, then the Dask cluster should be started outside of the Jupyter notebook in order to keep working when the kernel stops. We refer the user to the [dask documentation](http://distributed.dask.org/en/stable/) and to [this example](https://pangeo.io/setup_guides/hpc.html) for reference in using dask on an HPC system.<jupyter_code>store.simulate(wait_for_results=False)<jupyter_output><empty_output><jupyter_text>and when the simulations are complete, the (interrupted) workflow can continue as<jupyter_code>rejoinedstore = swyft.Store.load(cwd+'/mystore', simulator=simulator)
dataset = swyft.Dataset(Ntrain*8/10, prior, rejoinedstore, bound = new_bound)
post = swyft.Posteriors(dataset)
post.add([(0, 1)], device=DEVICE)
post.train([(0, 1)], max_epochs = 20)<jupyter_output>Training: lr=0.001, Epoch=20, VL=0.1976<jupyter_text>Optionally continuing iteratively by defining a new truncated prior. Finally, the posterior can be sampled<jupyter_code>samples = post.sample(1000000, obs0)
swyft.plot_corner(samples, [0, 1], color='k', figsize = (8,8), truth=v0, bins = 40);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08699730331809122, 'fraction_numerical': 0.011959198030249736, 'mean_word_length': 5.0240112994350286, 'pattern_counts': {'":': 2, '<': 59, '<?xml version=': 0, '>': 59, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '16', 'dir': 'jupyter-structured-clean-dedup', 'id': '114333', 'license': 'MIT', 'n_tokens_mistral': 2614, 'n_tokens_neox': 2484, 'n_words': 953, 'path': 'notebooks/Examples - 6. Dask Simulator.ipynb', 'repo_name': 'undark-lab/swyft', 'size': '8529'}