Repository: SchriderLab/Timesweeper (language: python)

Issue:
Added additional fisher's exact test functionality and cnn hyperparameter tuning code
Added functions that runs fisher's exact test over the entire dataset and a new script that does hyperparameter tuning for CNN.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 68412bd..8cb6c8c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,3 +12,4 @@ dist/
 *vcf
 *.tbi
 examples/*/
+SF2
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
deleted file mode 100644
index a56d03c..0000000
--- a/Dockerfile
+++ /dev/null
@@ -1,9 +0,0 @@
-FROM continuumio/miniconda3
-LABEL org.opencontainers.image.authors="Logan Whitehouse; lswhiteh@unc.edu, Dan Schrider; drs@unc.edu"
-
-COPY . $HOME/Timesweeper
-RUN conda install -c conda-forge mamba -y \
-    && mamba env create -f Timesweeper/blinx.yaml
-
-WORKDIR $HOME/Timesweeper
-ENTRYPOINT ["conda", "run", "-n", "blinx", "timesweeper"]
\ No newline at end of file
diff --git a/Makefile b/Makefile
deleted file mode 100644
index 7b745bd..0000000
--- a/Makefile
+++ /dev/null
@@ -1,11 +0,0 @@
-.PHONY: env
-
-all: env build
-
-env:
-	conda env create -f blinx.yml
-
-build:
-	rm src/timesweeper/timesweeper.egg-info/ -rf
-	python -m build
-	conda activate blinx; pip install -e .
\ No newline at end of file
diff --git a/README.md b/README.md
index 8a7dcfb..c1d5957 100644
--- a/README.md
+++ b/README.md
@@ -10,19 +10,21 @@ Experiments and figures for the Timesweeper manuscript can be found here: https:
   - [Table of Contents](#table-of-contents)
   - [Overview](#overview)
   - [Workflow Overview](#workflow-overview)
-  - [Installation](#installation)
+  - [Installation and Requirements](#installation-and-requirements)
   - [Timesweeper Configuration](#timesweeper-configuration)
     - [Configs required for both types of simulation:](#configs-required-for-both-types-of-simulation)
     - [Additional configs needed for stdpopsim simulation:](#additional-configs-needed-for-stdpopsim-simulation)
   - [Timesweeper Modules Details](#timesweeper-modules-details)
+    - [Overview](#overview-1)
     - [Custom Simulation (`simulate_custom`)](#custom-simulation-simulate_custom)
     - [stpopsim Simulation (`simulate_stdpopsim`)](#stpopsim-simulation-simulate_stdpopsim)
     - [Process VCF Files (`process`)](#process-vcf-files-process)
     - [Make Training Data (`condense`)](#make-training-data-condense)
     - [Neural Networks (`train`)](#neural-networks-train)
     - [Detect Sweeps (`detect`)](#detect-sweeps-detect)
-  - [Preparing Input Data for Timesweeper](#preparing-input-data-for-timesweeper)
+  - [Preparing Input Data for Timesweeper to Predict On](#preparing-input-data-for-timesweeper-to-predict-on)
   - [Example Usage](#example-usage)
+  - [Using Non-Timesweeper Simulations](#using-non-timesweeper-simulations)
 
 ## Overview
 
@@ -36,156 +38,180 @@ Timesweeper is built as a series of modules that are chained together to build a
 2. Simulate demographic model with time-series sampling
    1. `sim_custom` if using custom SLiM script
    2. `sim_stdpopsim` if using a SLiM script output by stdpopsim
-   3. Note: If available, we suggest using a job submission platform such as SLURM to parallelize simulations. This is the most resource and time-intensive part of the module by far.
-3. Preprocess simulated vcfs by merging with `process`
+   - Note: If available, we suggest using a job submission platform such as SLURM to parallelize simulations. This is the most resource and time-intensive part of the module by far.
+   - Optional: Preprocess VCFs simulated without timesweepers simulation modules by merging with `process_vcfs`
 4. Create features for the neural network with `condense`
 5. Train networks with `train`
 6. Run `detect` on VCF of interest using trained models and input data
 
 ---
 
-## Installation
+## Installation and Requirements
 
-If you have GNU Make installed the Timesweeper conda environment and SLiM can be downloaded and setup using
+The only required utility not packaged with the PyPi installation of Timesweeper is SLiM, which can be easily installed with conda through the conda-forge channel (see below) or as a binary.
 
-```{bash}
-git clone git@github.com:SchriderLab/timeSeriesSweeps.git
-
-cd timeSeriesSweeps
-make
-```
-
-Otherwise you can install dependencies the long way with:
+Otherwise see either [setup.cfg](setup.cfg) or [requirements.txt](requirements.txt) for general and specific requirements respectively.
 
+Timesweeper and all requirements can be installed from pip, I recommend doing so inside a virtual environment along with SLiM for easy access when simulating:
 ```{bash}
-git clone git@github.com:SchriderLab/timeSeriesSweeps.git
-
-conda env create -f blinx.yml
-
+conda create -n blinx -c conda-forge python slim
 conda activate blinx
-
-pip install .
+pip install timesweeper
 ```
 
 ---
 
 ## Timesweeper Configuration
 
-For any given experiment run you will need a YAML configuration file (see `example_timesweeper_config.yaml` for template) or the patience to write out all of your configuration parameters as arguments to each step of the workflow. The requirements for configuration are pretty different between using stdpopsim and custom SLiM script simulations, so I'll lay them out separately here. Each config argument will have the easy to read name and then will be followed by the YAML identifier to use in the config file.
+For any given experiment run you will need a YAML configuration file (see `example_timesweeper_config.yaml` for template). The requirements for configuration are different between using stdpopsim and custom SLiM script simulations.
 
 ### Configs required for both types of simulation:
 
+Example config file for a custom simulation run:
+
+```
+work dir: win_size_sims/
+slimfile: /slimfiles/onePop-selectiveSweep-vcf-noselcoeff.slim
+slim path: slim
+mut types: [2]
+scenarios: ["neut", "sdn", "ssv"]
+win_size: 51
+num_sample_points : 20
+inds_per_tp : 10  # Diploid inds
+sample sizes: [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
+ploidy: 2
+physLen : 5000000
+reps: 30000
+```
+
+Example config file for a stdpopsim simulation run:
+
+```
+work dir: OoA
+slimfile: OoA3pop.slim
+pop: p2
+sample sizes: [10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
+years sampled: [4750, 4500, 4250, 4000, 3750, 3500, 3250, 3000, 2750, 2500, 2250, 2000, 1750, 1500, 1250, 1000, 750, 500, 250, 0]
+ploidy: 2
+scenarios: ["neut", "sdn", "ssv"]
+mut types: [2]
+win_size: 51
+inds_per_tp: 10
+num_sample_points : 20
+reps: 10000
+selection gen: 50 # Gens before first sampling
+selection coeff bounds: [0.00025, 0.25]
+mut rate: 1.29e-8
+slim path: slim
+gen time: 25
+```
+
 - **Working Directory** (`work dir`) - this serves as the "home base" of an experiment. This is where all intermediate and final data will be written to.
 - **SLiM Script** (`slimfile`) - either generated by stdpopsim (see "Using stdpopsim to generate a SLiM template" below) or from a custom SLiM script (see "Using a custom SLiM script").
+- - **Path to SLiM Executable** (`slim path`) - if you use the Makefile will be `/<path>/<to>/<timesweeper>/SLiM/build/slim`
+- **Mutation types** (`mut types`) - target mutation types output by SLiM in outputVCFSample. Typically MT=1 is for neutral mutations and MT=2 is for selected alleles.
+- **Scenarios** (`scenarios`) - Names of the types of simulation scenarios input into Timesweeper. The standard (and maybe only?) option for current implementation is `["neut", "ssv", "sdn"]`.
+- **Window size** (`win_size`) - Number of SNPs to use to create feature vectors.
+- **Number of Sample Points** (`num_sample_points`) - Number of timepoints sampled, identical to `len(sample_sizes)`.
+- **Individuals per timepoint** (`inds_per_tp`) - Another lazy variable to do some quick calculations, variable sampling sizes is supported but this is used for some quick calculations. Will probably be removed at some point when it causes errors.
 - **Sample sizes of each timepoint** (`sample sizes`) - sampled in your data (i.e. how many individuals are included in timepoint 1, timepoint 2, etc)
-- **Ploidy** of your samples (yes, believe it or not we support not just diploids!)
-- **Simulation Replicates** (`reps`) - for each scenario: neutral, selection on de novo mutation, selection on standing variation. We like 2500, but you could probably get away with fewer!  
-- **Path to SLiM Executable** (`slim path`) - if you use the Makefile will be `/<path>/<to>/<timesweeper>/SLiM/build/slim`
+- **Ploidy** (`ploidy`) - ploidy of your samples.
+- **Physical size** (`physLen`) - Size of the chromosome to simulate. Will be overwritten by stdpopsim if used.
+- **Simulation Replicates** (`reps`) - for each scenario: neutral, selection on de novo mutation, selection on standing variation. Will be overwritten with `--rep-range` argument if doing parallelized sims.
+
 
 ### Additional configs needed for stdpopsim simulation: 
 - **Target Population ID** (`pop`) - will be something like "p1" or "p2", will be identified in the stdpopsim model catalog.
 - **Years Sampled** (`years sampled`) - typically in years Before Present (BP - which is years prior to 1950 apparently). This is only needed if you are injecting sampling into a stdpopsim model or if you'd like to calculate FIT values along with the AFT predictions with Timesweeper. If you don't have any idea what these values should be you can still run Timesweeper and get AFT predictions no problem, just use a custom SLiM script and remove this line from the config file.
 - **Selection Generation Before Sampling** (`selection gen`) - Number of generations before the __first__ sample timepoint. We've found that you can be relatively flexible (see the manuscript), but 50 is a good default value.
-- **Selection Coefficient Bounds** (`selection coeff bounds`) - to improve robustness we use a log-normal distribution to draw selection coefficients with lower and upper bounds specified here. If you want a non-stochastic selection coefficient simply use the same number twice.
+- **Selection Coefficient Bounds** (`selection coeff bounds`) - to improve robustness we use a uniform distribution to draw selection coefficients with lower and upper bounds specified here. If you want a non-stochastic selection coefficient simply use the same number twice.
 - **Mutation Rate** (`mut rate`) - just overwrites the stdpopsim mutation rate in case you'd like to fiddle with it. 
 - **Generation Time** (`gen time`) - allows conversions between generations and continuous time. 
+- **Selection Generation** (`selection gen`) - Number of generations to start sampling (with a random offset each replicate) before selection.
 
-Example config file for a stdpopsim simulation run:
 
-```{yaml}
-#General
-work dir: example_experiment
-slimfile: example_model.slim
-pop: p1
-sample sizes: [5, 4, 6, 4, 5]
-years sampled: [500, 400, 300, 200, 100]
-ploidy: 2
+---
 
-#Simulation
-reps: 2500
-selection gen: 200 # Gens before first sampling
-selection coeff bounds: [0.005, 0.5]
-mut rate: 1.29e-8
-gen time: 25
+## Timesweeper Modules Details
+
+### Overview
 ```
+$ timesweeper
+usage: timesweeper [-h] {sim_stdpopsim,sim_custom,process,condense,train,detect,plot_training,plot_freqs,summarize,merge_logs} ...
 
----
+Timesweeper CLI
 
-## Timesweeper Modules Details
+positional arguments:
+  {sim_stdpopsim,sim_custom,process,condense,train,detect,plot_training,plot_freqs,summarize,merge_logs}
+                        Timesweeper modes
+    sim_stdpopsim       Injects time-series sampling into stdpopsim SLiM output script.
+    sim_custom          Simulates selection for training Timesweeper using a pre-made SLiM script.
+    process             Module for splitting multivcfs (vertically concatenated vcfs) into merged (horizontally concatenated vcfs) if simulating
+                        without the sim module.
+    condense            Creates training data from simulated merged vcfs after process_vcfs.py has been run.
+    train               Handler script for neural network training and prediction for TimeSweeper Package. Will train two models: one for the series
+                        of timepoints generated using the hfs vectors over a timepoint and one
+    detect              Module for iterating across windows in a time-series vcf file and predicting whether a sweep is present at each snp-
+                        centralized window.
+    plot_training       Plots central SNPs from simulations to visually inspect mean trends over replicates.
+    plot_freqs          Create a bedfile of major and minor allele frequency changes over time.
+    summarize           Creates a CSV of data parsed from slim log files.
+    merge_logs          Merges the summary TSV from SLiM logs with test data predictions.
+
+options:
+  -h, --help            show this help message and exit
+```
 
 ### Custom Simulation (`simulate_custom`) 
 
 A flexible wrapper for a SLiM script that assumes you have a demographic model already defined in SLiM and would like to use it for Timesweeper. This module is meant to be modified with utility functions and other constants to feed into simulation replicates using the string variable `d_block`. This set of constants will be fed to SLiM at runtime but can be modified in the module as needed prior. The module on GitHub is the version used to generate simulations for the manuscript, and is meant to provide ideas/context for how it's possible to modify the `d_block` and write utility functions for it.
 
-   There are some basic requirements for using a custom SLiM script:
-   1. Each timepoint must be output using a the outputVCFSample function like so: `<pop_id>.outputVCFSample(<num_individuals>, replace=T, filePath=outFile, append=T);`
-   2. The constants `sweep`, `outFile`, and `dumpFile` must be used in your simulation script
+  There are some basic requirements for using a custom SLiM script:
+  1. Each timepoint must be output using a the outputVCFSample function like so: `<pop_id>.outputVCFSample(<num_individuals>, replace=T, filePath=outFile, append=T);`
+  2. The constants `sweep`, `outFile`, and `dumpFile` must be used in your simulation script
       - `sweep`: one of "neut"/"sdn"/"soft", equivalent to neutral, selection on *de novo* mutation, and selection on standing variation respectively. This identifier is used both in the SLiM script to condition on scenarios but also in the output file naming.
       - `outFile`: is the VCF file that will be used as output for samples for a given replicate. Will be set as `<work_dir>/vcfs/<sweep>/<rep>.multivcf`
       - `dumpFile`: similarly to outFile this is where the intermediate simulation state is saved to in case of mutation loss or other problems with a replicate.
 
+  There is an example SLiM simulation in [constant_population.slim](timesweeper/constant_population.slim) that has been used and templated/modified for many experiments for the manuscript and is ready to be used with the current `d_block` setup in `sim_custom.py`.
+
 ```
 $ timesweeper sim_custom -h
 usage: timesweeper sim_custom [-h] [--threads THREADS]
-                          [--rep-range REP_RANGE REP_RANGE]
-                          {yaml,cli} ...
+                              [--rep-range REP_RANGE REP_RANGE] -y YAML_CONFIG
 
 Simulates selection for training Timesweeper using a pre-made SLiM script.
 
-positional arguments:
-  {yaml,cli}
-
-optional arguments:
+options:
   -h, --help            show this help message and exit
-  --threads THREADS     Number of processes to parallelize across. Defaults to all.
-  --rep-range START STOP
+  --threads THREADS     Number of processes to parallelize across. Defaults to
+                        all.
+  --rep-range REP_RANGE REP_RANGE
                         <start, stop>. If used, only range(start, stop) will
                         be simulated for reps. This is to allow for easy SLURM
                         parallel simulations.
+  -y YAML_CONFIG, --yaml YAML_CONFIG
+                        YAML config file with all required options defined.
 
-$ timesweeper sim_custom cli -h
-usage: timesweeper sim_custom cli [-h] [-w WORK_DIR] -i SLIM_FILE
-                              [--slim-path SLIM_PATH] [--reps REPS]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -w WORK_DIR, --work-dir WORK_DIR
-                        Directory used as work dir for simulate modules.
-                        Should contain simulated vcfs processed using
-                        process_vcf.py.
-  -i SLIM_FILE, --slim-file SLIM_FILE
-                        SLiM Script to simulate with. Must output to a single
-                        VCF file.
-  --slim-path SLIM_PATH
-                        Path to SLiM executable.
-  --reps REPS           Number of replicate simulations to run if not using rep-range.
-
-timesweeper sim_custom yaml -h
-usage: timesweeper sim_custom yaml [-h] YAML_CONFIG
-
-positional arguments:
-  YAML_CONFIG  YAML config file with all cli options defined.
-
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 
 ### stpopsim Simulation (`simulate_stdpopsim`) 
 
-For use with SLiM scripts that have been generated using stdpopsim's `--slim-script` option to output the model. This allows for out of the box demographic models downloaded straight from the catalog stdpopsim adds to regularly. Some information needs to be gotten from the model definition so that the wrapper knows which population to sample from, how to scale values if rescaling the simulation, and more. These are described in detail both in the help message of the module and in the above doc section "Configs required for both types of simulation".
+For use with SLiM scripts that have been generated using stdpopsim's `--slim-script` option to output the model. This allows for out of the box demographic models downloaded straight from the catalog stdpopsim adds to regularly. Some information needs to be gotten from the model definition so that the wrapper knows which population to sample from, how to scale values if rescaling the simulation, and more. These are described in detail both in the help message of the module and in the above doc section [Configs required for both types of simulation](#configs-required-for-both-types-of-simulation).
 
-```$ timesweeper sim_stdpopsim -h
+```
+$ timesweeper sim_stdpopsim -h
 usage: timesweeper sim_stdpopsim [-h] [-v] [--threads THREADS]
-                             [--rep-range REP_RANGE REP_RANGE]
-                             {yaml,cli} ...
+                                 [--rep-range REP_RANGE REP_RANGE]
+                                 YAML CONFIG
 
 Injects time-series sampling into stdpopsim SLiM output script.
 
 positional arguments:
-  {yaml,cli}
+  YAML CONFIG           YAML config file with all options defined.
 
-optional arguments:
+options:
   -h, --help            show this help message and exit
   -v, --verbose         Print verbose logging during process.
   --threads THREADS     Number of processes to parallelize across.
@@ -193,107 +219,28 @@ optional arguments:
                         <start, stop>. If used, only range(start, stop) will
                         be simulated for reps. This is to allow for easy SLURM
                         parallel simulations.
-
-timesweeper sim_stdpopsim cli -h
-usage: timesweeper sim_stdpopsim cli [-h] -i SLIM_FILE --reps REPS [--pop POP]
-                                 --sample_sizes SAMPLE_SIZES
-                                 [SAMPLE_SIZES ...] --years-sampled
-                                 YEARS_SAMPLED [YEARS_SAMPLED ...]
-                                 [--selection-generation SEL_GEN]
-                                 [--selection-coeff-bounds SEL_COEFF_BOUNDS SEL_COEFF_BOUNDS]
-                                 [--mut-rate MUT_RATE] [--work-dir WORK_DIR]
-                                 [--slim-path SLIM_PATH]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -i SLIM_FILE, --slim-file SLIM_FILE
-                        SLiM Script output by stdpopsim to add time-series
-                        sampling to.
-  --reps REPS           Number of replicate simulations to run. If using
-                        rep_range can just fill with random int.
-  --pop POP             Label of population to sample from, will be defined in
-                        SLiM Script. Defaults to p2.
-  --sample_sizes SAMPLE_SIZES [SAMPLE_SIZES ...]
-                        Number of individuals to sample without replacement at
-                        each sampling point. Will be multiplied by ploidy to
-                        sample chroms from slim. Must match the number of
-                        entries in the -y flag.
-  --years-sampled YEARS_SAMPLED [YEARS_SAMPLED ...]
-                        Years BP (before 1950) that samples are estimated to
-                        be from. Must match the number of entries in the --sample-sizes
-                        flag.
-  --selection-generation SEL_GEN
-                        Number of gens before first sampling to introduce
-                        selection in population. Defaults to 200.
-  --selection-coeff-bounds SEL_COEFF_BOUNDS SEL_COEFF_BOUNDS
-                        Bounds of log-uniform distribution for pulling
-                        selection coefficient of mutation being introduced.
-                        Defaults to [0.005, 0.5]
-  --mut-rate MUT_RATE   Mutation rate for mutations not being tracked for
-                        sweep detection. Defaults to 1.29e-8 as defined in
-                        stdpopsim for OoA model.
-  --work-dir WORK_DIR   Directory to start workflow in, subdirs will be
-                        created to write simulation files to. Will be used in
-                        downstream processing as well.
-  --slim-path SLIM_PATH
-                        Path to SLiM executable.
-
-$ timesweeper sim_stdpopsim yaml -h
-usage: timesweeper sim_stdpopsim yaml [-h] YAML CONFIG
-
-positional arguments:
-  YAML CONFIG  YAML config file with all cli options defined.
-
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 ### Process VCF Files (`process`) 
 
-This module splits the multivcf files (which are just multiple concatenated VCF entries) generated by SLiM and then merges them in order from most ancient to most current timepoints. This replicates the preparation process needed for Timesweeper input data but at scale and automatically parameterized for data simulated using Timesweeper's simulation modules.
-
+This module splits the multivcf files (which are just multiple concatenated VCF entries) generated by SLiM and then merges them in order from most ancient to most current timepoints. This is already integrated into `sim_custom` and `sim_stdpopsim` but can be used separately if you would like to simulate without the sim_custom or sim_stdpopsim modules.
 
 ```
 $ timesweeper process -h
-usage: timesweeper process [-h] [--vcf-header VCF_HEADER] [--threads THREADS]
-                       {yaml,cli} ...
+usage: timesweeper process [-h] -i IN_DIR -t THREADS
 
-Splits and re-merges VCF files to prepare for fast feature creation.
+Module for splitting multivcfs (vertically concatenated vcfs) into merged
+(horizontally concatenated vcfs) if simulating without the sim module.
 
-positional arguments:
-  {yaml,cli}
-
-optional arguments:
+options:
   -h, --help            show this help message and exit
-  --vcf-header VCF_HEADER
-                        String that tops VCF header, used to split entries to
-                        new files.
-  --threads THREADS     Number of processes to parallelize across.
+  -i IN_DIR, --in-dir IN_DIR
+                        Top-level directory containing subidrectories labeled
+                        with the names of each scenario and then replicate
+                        numbers inside those containing each multivcf.
+  -t THREADS, --threads THREADS
+                        Threads to use for multiprocessing.
 
-$ timesweeper process cli -h
-usage: timesweeper process cli [-h] [-w WORK_DIR] --sample_sizes SAMPLE_SIZES
-                           [SAMPLE_SIZES ...]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -w WORK_DIR, --work-dir WORK_DIR
-                        Directory used as work dir for simulate modules.
-                        Should contain simulated vcfs processed using
-                        process_vcf.py.
-  --sample_sizes SAMPLE_SIZES [SAMPLE_SIZES ...]
-                        Number of individuals to sample without replacement at
-                        each sampling point. Will be multiplied by ploidy to
-                        sample chroms from slim. Must match the number of
-                        entries in the -y flag.
-
-$ timesweeper process yaml -h
-usage: timesweeper process yaml [-h] YAML CONFIG
-
-positional arguments:
-  YAML CONFIG  YAML config file with all cli options defined.
-
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 ### Make Training Data (`condense`) 
@@ -306,46 +253,56 @@ Note: the process of retrieving known-selection sites is based on the mutation t
 
 ```
 $ timesweeper condense -h
-usage: timesweeper condense [-h] [--threads THREADS] [-m MISSINGNESS]
-                                 {yaml,cli} ...
+usage: timesweeper condense [-h] [--threads THREADS] [-o OUTFILE]
+                            [--subsample-inds SUBSAMPLE_INDS]
+                            [--subsample-tps SUBSAMPLE_TPS] [--og-tps OG_TPS]
+                            [-m MISSINGNESS] [-f FREQ_INC_THRESHOLD] [-a]
+                            [--hft] [--no-progress] [--verbose] -y YAML_CONFIG
 
-Creates training data from simulated merged vcfs after timesweeper process has
+Creates training data from simulated merged vcfs after process_vcfs.py has
 been run.
 
-positional arguments:
-  {yaml,cli}
-
-optional arguments:
+options:
   -h, --help            show this help message and exit
   --threads THREADS     Number of processes to parallelize across.
+  -o OUTFILE, --outfile OUTFILE
+                        Pickle file to dump dictionaries with training data
+                        to. Should probably end with .pkl.
+  --subsample-inds SUBSAMPLE_INDS
+                        Number of individuals to subsample if using a larger
+                        simulation than needed for efficiency. NOTE: If you
+                        use this, keep the sample_sizes entry in the yaml
+                        config identical to the original simulations. This is
+                        mostly for the paper experiments and as such isn't
+                        very cleanly implemented.
+  --subsample-tps SUBSAMPLE_TPS
+                        Number of timepoints to subsample if using a larger
+                        simulation than needed for efficiency. NOTE: If you
+                        use this, keep the sample_sizes entry in the yaml
+                        config identical to the original simulations. This is
+                        mostly for the paper experiments and as such isn't
+                        very cleanly implemented.
+  --og-tps OG_TPS       Number of timepoints taken in the original data to be
+                        subsetted if using --subample-tps.
   -m MISSINGNESS, --missingness MISSINGNESS
                         Missingness rate in range of [0,1], used as the
                         parameter of a binomial distribution for randomly
                         removing known values.
+  -f FREQ_INC_THRESHOLD, --freq-increase-threshold FREQ_INC_THRESHOLD
+                        If given, only include sim replicates where the sweep
+                        site has a minimum increase of <freq_inc_thr> from the
+                        first timepoint to the last.
+  -a, --allow-shoulders
+                        1/3 of samples from sweep classes will be offset to be
+                        used as neutral shoulders.
+  --hft                 Whether to calculate HFT alongside AFT.
+                        Computationally more expensive.
+  --no-progress         Turn off progress bar.
+  --verbose             Raise warnings from issues usually stemming from bad
+                        replicates.
+  -y YAML_CONFIG, --yaml YAML_CONFIG
+                        YAML config file with all required options defined.
 
-$ timesweeper condense cli -h
-usage: timesweeper condense cli [-h] [-w WORK_DIR] -s SAMP_SIZES
-                                     [SAMP_SIZES ...]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -w WORK_DIR, --work-dir WORK_DIR
-                        Directory used as work dir for simulate modules.
-                        Should contain simulated vcfs processed using
-                        process_vcf.py.
-  -s SAMP_SIZES [SAMP_SIZES ...], --sample-sizes SAMP_SIZES [SAMP_SIZES ...]
-                        Number of individuals from each timepoint sampled.
-                        Used to index VCF data from earliest to latest
-                        sampling points.
-
-$ timesweeper condense yaml -h
-usage: timesweeper condense yaml [-h] YAML CONFIG
-
-positional arguments:
-  YAML CONFIG  YAML config file with all cli options defined.
-
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 ### Neural Networks (`train`) 
@@ -354,72 +311,63 @@ Timesweeper's neural network architecture is a shallow 1DCNN implemented in Kera
 
 ```
 $ timesweeper train -h
-usage: timesweeper train [-h] [-n EXPERIMENT_NAME] {yaml,cli} ...
+usage: timesweeper train [-h] -i TRAINING_DATA -d DATA_TYPE
+                         [-s SUBSAMPLE_AMOUNT] [-n EXPERIMENT_NAME] -y
+                         YAML_CONFIG [--single-tp]
 
 Handler script for neural network training and prediction for TimeSweeper
 Package. Will train two models: one for the series of timepoints generated
 using the hfs vectors over a timepoint and one
 
-positional arguments:
-  {yaml,cli}
-
-optional arguments:
+options:
   -h, --help            show this help message and exit
+  -i TRAINING_DATA, --training-data TRAINING_DATA
+                        Pickle file containing data formatted with
+                        make_training_features.py.
+  -d DATA_TYPE, --data-type DATA_TYPE
+                        AFT or HFT data preparation.
+  -s SUBSAMPLE_AMOUNT, --subsample-amount SUBSAMPLE_AMOUNT
+                        Amount of data to subsample for each class to test for
+                        sample size effects.
   -n EXPERIMENT_NAME, --experiment-name EXPERIMENT_NAME
                         Identifier for the experiment used to generate the
                         data. Optional, but helpful in differentiating runs.
+  -y YAML_CONFIG, --yaml YAML_CONFIG
+                        YAML config file with all required options defined.
+  --single-tp           Whether to use the tp1_model module for a special case
+                        experiment.
 
-$ timesweeper train cli -h
-usage: timesweeper train cli [-h] [-w WORK_DIR]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -w WORK_DIR, --work-dir WORK_DIR
-                        Directory used as work dir for simulate modules.
-                        Should contain pickled training data from simulated vcfs processed using
-                        process_vcf.py.
-
-$ timesweeper train yaml -h
-usage: timesweeper train yaml [-h] YAML CONFIG
-
-positional arguments:
-  YAML CONFIG  YAML config file with all cli options defined.
-
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 ### Detect Sweeps (`detect`) 
 
 Finally, the main module of the package is for detecting sweeps in a given VCF. This loads in the prepared VCF (see "Preparing Input Data for Timesweeper" below) in chunks, converts allele data to time-series allele velocity data, and predicts using the 1DCNN trained on simulated data. Each prediction represents a 51-SNP window with the focal allele being the actual target.
 
-Timesweeper outputs predictions as both a csv file with headers `Chrom	BP	Class	Neut Score	Hard Score	Soft Score	Win_Start	Win_End` and a BED file with the headers `Chrom	Win_Start	Win_End	BP`. The BED file allows for easy intersections using bedtools and can be cross-referenced back to the CSV for score filtering.
+Timesweeper outputs predictions as both a csv file and a bedfile. The BED file allows for easy intersections using bedtools and can be cross-referenced back to the CSV for score filtering.
 
 Here are the details on the headers:
 - Chrom: Chromosome/contig, identical to VCF file name of it
 - BP: location of central allele in the window being predicted on
 - Class: Neut/Hard/Soft, class identified by the maximum score in 3-class softmax output from the model
-- Neut/Hard/Soft Score: raw score from softmax final layer of 1DCNN
+- Neut/Hard/Soft Prob: raw score from softmax final layer of 1DCNN
+- SDN/SSV_sval: Predicted selection coefficient by each (ssv, sdn) regression model
 - Win_Start/End: left and right-most locations of the SNPs on each side of the window being predicted on
 
 Note: By default Timesweeper only outputs sites with a minimum sweep (sdn+soft) score of 0.66 to prevent massive amounts of neutral outputs. This value could easily be modified in the module but we find it better to filter after the fact for more flexibility.
 
-Timesweeper will optionally run frequency increment test if the generation time and years sampled are provided. In some cases this may not be available/desired, in which case if either of those values are left out it will not be run.
-
 Timesweeper also has a `--benchmark` flag that will allow for testing accuracy on simulated data if wanted. This will search the input data for the mutation type identifier flags allowing a benchmark of detection accuracy on data that has a ground truth.
 
 ```
 $ timesweeper detect -h
-usage: timesweeper detect [-h] -i INPUT_VCF [--benchmark] --aft-model AFT_MODEL
-                      {yaml,cli} ...
+usage: timesweeper detect [-h] -i INPUT_VCF [--benchmark] --aft-class-model
+                          AFT_CLASS_MODEL --aft-reg-model AFT_REG_MODEL
+                          [--hft-class-model HFT_CLASS_MODEL] --hft-reg-model
+                          HFT_REG_MODEL -o OUTFILE -y YAML_CONFIG -s SCALAR
 
 Module for iterating across windows in a time-series vcf file and predicting
 whether a sweep is present at each snp-centralized window.
 
-positional arguments:
-  {yaml,cli}
-
-optional arguments:
+options:
   -h, --help            show this help message and exit
   -i INPUT_VCF, --input-vcf INPUT_VCF
                         Merged VCF to scan for sweeps. Must be merged VCF
@@ -430,49 +378,36 @@ optional arguments:
                         outputVCFSample as well as neutral predictions, use
                         this flag. Otherwise the mutation type will not be
                         looked for in the VCF entry nor reported with results.
-  --aft-model AFT_MODEL
-                        Path to Keras2-style saved model to load for aft
-                        prediction.
-
-$ timesweeper detect cli -h
-usage: timesweeper detect cli [-h] -s SAMP_SIZES [SAMP_SIZES ...] [-w WORKING_DIR]
-                          [--years-sampled YEARS_SAMPLED [YEARS_SAMPLED ...]]
-                          [--gen-time GEN_TIME]
-
-optional arguments:
-  -h, --help            show this help message and exit
-  -s SAMP_SIZES [SAMP_SIZES ...], --sample-sizes SAMP_SIZES [SAMP_SIZES ...]
-                        Number of individuals from each timepoint sampled.
-                        Used to index VCF data from earliest to latest
-                        sampling points.
-  -w WORKING_DIR, --work-dir WORKING_DIR
-                        Working directory for workflow, should be identical to
-                        previous steps.
-  --years-sampled YEARS_SAMPLED [YEARS_SAMPLED ...]
-                        Years BP (before 1950) that samples are estimated to
-                        be from. Only used for FIT calculations, and is
-                        optional if you don't care about those.
-  --gen-time GEN_TIME   Generation time to multiply years_sampled by.
-                        Similarly to years_sampled, only used for FIT
-                        calculation and is optional.
-
-$ timesweeper detect yaml -h
-usage: timesweeper detect yaml [-h] YAML CONFIG
-
-positional arguments:
-  YAML CONFIG  YAML config file with all cli options defined.
+  --aft-class-model AFT_CLASS_MODEL
+                        Path to Keras2-style saved model to load for
+                        classification aft prediction.
+  --aft-reg-model AFT_REG_MODEL
+                        Path to Keras2-style saved model to load for
+                        regression aft prediction. Either SDN or SSV work,
+                        both will be loaded.
+  --hft-class-model HFT_CLASS_MODEL
+                        Path to Keras2-style saved model to load for
+                        classification hft prediction.
+  --hft-reg-model HFT_REG_MODEL
+                        Path to Keras2-style saved model to load for
+                        regression hft prediction.
+  -o OUTFILE, --outfile-prefix OUTFILE
+                        Prefix to use for results csvs. Will have
+                        '_[aft/hft].csv' appended to it.
+  -y YAML_CONFIG, --yaml YAML_CONFIG
+                        YAML config file with all required options defined.
+  -s SCALAR, --scalar SCALAR
+                        Minmax scalar saved as a pickle file during training.
 
-optional arguments:
-  -h, --help   show this help message and exit
 ```
 
 ---
 
-## Preparing Input Data for Timesweeper
+## Preparing Input Data for Timesweeper to Predict On
 
-Timesweeper needs a specific format of input data to work, namely a VCF file merged to the superset of all samples and polymorphisms that is merged *in order* of oldest to most recent. That last point is extremely important and easily missed
+Timesweeper needs a specific format of input data to work, namely a VCF file merged to the superset of all samples and polymorphisms that is merged *in order* of oldest to most recent. That last point is extremely important and easily missed.
 
-VCFs of all samples will need to be merged using the `bcftools merge -Oz --force-samples <inputs_earliest.vcf ... inputs_latest.vcf> > merged.vcf.gz` options.
+VCFs of all samples will need to be merged using the `bcftools merge -Oz --force-samples -0 <inputs_earliest.vcf ... inputs_latest.vcf> > merged.vcf.gz` options.
 
 ---
 
@@ -485,18 +420,22 @@ conda activate blinx
 cd timesweeper
 
 #Simulate training data
-timesweeper sim_custom yaml example_config.yaml
+timesweeper sim_custom yaml examples/example_config.yaml
 
-#Process VCFs 
-timesweeper process yaml example_config.yaml
-
-#Assume foo.vcf has a missingness of 0.05 and create pickle file
-timesweeper condense -m 0.05 yaml example_config.yaml
+#Create feature vectors for both AFT and HFT
+timesweeper condense --hft yaml examples/example_config.yaml
 
 #Train network
-timesweeper train -n example_ts_run yaml example_config.yaml
+timesweeper train -n example_ts_run yaml examples/example_config.yaml
 
 #Predict on input VCF
-timesweeper detect -i foo.vcf --aft-model ts_experiment/trained_models/example_ts_run_Timesweeper_aft
+timesweeper detect -i foo.vcf \
+  --aft-class-model ts_experiment/trained_models/example_ts_run_Timesweeper_aft \
+  --hft-class-model ts_experiment/trained_models/example_ts_run_Timesweeper_hft \
+  --aft-reg-model ts_experiment/trained_models/REG_example_ts_run_Timesweeper_aft_sdn \
+  --hft-reg-model ts_experiment/trained_models/REG_example_ts_run_Timesweeper_hft_sdn \
+  
 ```
 
+
+## Using Non-Timesweeper Simulations
\ No newline at end of file
diff --git a/blinx.yaml b/blinx.yaml
deleted file mode 100644
index 2989e50..0000000
--- a/blinx.yaml
+++ /dev/null
@@ -1,251 +0,0 @@
-name: blinx
-channels:
-  - conda-forge
-  - bioconda
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1=conda_forge
-  - _openmp_mutex=4.5=2_gnu
-  - absl-py=1.2.0=pyhd8ed1ab_0
-  - aiohttp=3.8.3=py310h5764c6d_0
-  - aiosignal=1.2.0=pyhd8ed1ab_0
-  - alsa-lib=1.2.7.2=h166bdaf_0
-  - asciitree=0.3.3=py_2
-  - astunparse=1.6.3=pyhd8ed1ab_0
-  - async-timeout=4.0.2=pyhd8ed1ab_0
-  - attr=2.5.1=h166bdaf_1
-  - attrs=22.1.0=pyh71513ae_1
-  - bcftools=1.16=hfe4b78e_0
-  - blinker=1.5=pyhd8ed1ab_0
-  - bokeh=2.4.3=pyhd8ed1ab_3
-  - brotli=1.0.9=h166bdaf_7
-  - brotli-bin=1.0.9=h166bdaf_7
-  - brotlipy=0.7.0=py310h5764c6d_1004
-  - bzip2=1.0.8=h7f98852_4
-  - c-ares=1.18.1=h7f98852_0
-  - ca-certificates=2022.9.24=ha878542_0
-  - cached-property=1.5.2=hd8ed1ab_1
-  - cached_property=1.5.2=pyha770c72_1
-  - cachetools=5.2.0=pyhd8ed1ab_0
-  - certifi=2022.9.24=pyhd8ed1ab_0
-  - cffi=1.15.1=py310h255011f_0
-  - charset-normalizer=2.1.1=pyhd8ed1ab_0
-  - click=8.1.3=py310hff52083_0
-  - cloudpickle=2.2.0=pyhd8ed1ab_0
-  - colorama=0.4.5=pyhd8ed1ab_0
-  - contourpy=1.0.5=py310hbf28c38_0
-  - cryptography=38.0.1=py310h597c629_0
-  - cycler=0.11.0=pyhd8ed1ab_0
-  - cytoolz=0.12.0=py310h5764c6d_0
-  - dask=2022.9.2=pyhd8ed1ab_0
-  - dask-core=2022.9.2=pyhd8ed1ab_0
-  - dbus=1.13.6=h5008d03_3
-  - distributed=2022.9.2=pyhd8ed1ab_0
-  - entrypoints=0.4=pyhd8ed1ab_0
-  - expat=2.4.9=h27087fc_0
-  - fasteners=0.17.3=pyhd8ed1ab_0
-  - fftw=3.3.10=nompi_hf0379b8_105
-  - flatbuffers=2.0.7=h27087fc_0
-  - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
-  - font-ttf-inconsolata=3.000=h77eed37_0
-  - font-ttf-source-code-pro=2.038=h77eed37_0
-  - font-ttf-ubuntu=0.83=hab24e00_0
-  - fontconfig=2.14.0=hc2a2eb6_1
-  - fonts-conda-ecosystem=1=0
-  - fonts-conda-forge=1=0
-  - fonttools=4.37.4=py310h5764c6d_0
-  - freetype=2.12.1=hca18f0e_0
-  - frozenlist=1.3.1=py310h5764c6d_0
-  - fsspec=2022.8.2=pyhd8ed1ab_0
-  - gast=0.4.0=pyh9f0ad1d_0
-  - gettext=0.19.8.1=h27087fc_1009
-  - giflib=5.2.1=h36c2ea0_2
-  - glib=2.74.0=h6239696_0
-  - glib-tools=2.74.0=h6239696_0
-  - google-auth=2.12.0=pyh1a96a4e_0
-  - google-auth-oauthlib=0.4.6=pyhd8ed1ab_0
-  - google-pasta=0.2.0=pyh8c360ce_0
-  - grpc-cpp=1.47.1=h05bd8bd_6
-  - grpcio=1.47.1=py310h5b20944_6
-  - gsl=2.7=he838d99_0
-  - gst-plugins-base=1.20.3=h57caac4_2
-  - gstreamer=1.20.3=hd4edc92_2
-  - h5py=3.7.0=nompi_py310h416281c_101
-  - hdf5=1.12.2=nompi_h2386368_100
-  - heapdict=1.0.1=py_0
-  - htslib=1.16=h6bc39ce_0
-  - icu=70.1=h27087fc_0
-  - idna=3.4=pyhd8ed1ab_0
-  - importlib-metadata=4.11.4=py310hff52083_0
-  - jack=1.9.18=h8c3723f_1003
-  - jinja2=3.1.2=pyhd8ed1ab_1
-  - joblib=1.2.0=pyhd8ed1ab_0
-  - jpeg=9e=h166bdaf_2
-  - keras=2.10.0=pyhd8ed1ab_0
-  - keras-preprocessing=1.1.2=pyhd8ed1ab_0
-  - keyutils=1.6.1=h166bdaf_0
-  - kiwisolver=1.4.4=py310hbf28c38_0
-  - krb5=1.19.3=h3790be6_0
-  - lcms2=2.12=hddcbb42_0
-  - ld_impl_linux-64=2.36.1=hea4e1c9_2
-  - lerc=4.0.0=h27087fc_0
-  - libabseil=20220623.0=cxx17_h48a1fff_4
-  - libblas=3.9.0=16_linux64_openblas
-  - libbrotlicommon=1.0.9=h166bdaf_7
-  - libbrotlidec=1.0.9=h166bdaf_7
-  - libbrotlienc=1.0.9=h166bdaf_7
-  - libcap=2.65=ha37c62d_0
-  - libcblas=3.9.0=16_linux64_openblas
-  - libclang=14.0.6=default_h2e3cab8_0
-  - libclang13=14.0.6=default_h3a83d3e_0
-  - libcups=2.3.3=h3e49a29_2
-  - libcurl=7.85.0=h7bff187_0
-  - libdb=6.2.32=h9c3ff4c_0
-  - libdeflate=1.13=h166bdaf_0
-  - libedit=3.1.20191231=he28a2e2_2
-  - libev=4.33=h516909a_1
-  - libevent=2.1.10=h9b69904_4
-  - libffi=3.4.2=h7f98852_5
-  - libflac=1.3.4=h27087fc_0
-  - libgcc-ng=12.1.0=h8d9b700_16
-  - libgfortran-ng=12.1.0=h69a702a_16
-  - libgfortran5=12.1.0=hdcd56e2_16
-  - libglib=2.74.0=h7a41b64_0
-  - libgomp=12.1.0=h8d9b700_16
-  - libiconv=1.17=h166bdaf_0
-  - liblapack=3.9.0=16_linux64_openblas
-  - libllvm14=14.0.6=he0ac6c6_0
-  - libnghttp2=1.47.0=hdcd2b5c_1
-  - libnsl=2.0.0=h7f98852_0
-  - libogg=1.3.4=h7f98852_1
-  - libopenblas=0.3.21=pthreads_h78a6416_3
-  - libopus=1.3.1=h7f98852_1
-  - libpng=1.6.38=h753d276_0
-  - libpq=14.5=hd77ab85_0
-  - libprotobuf=3.21.7=h6239696_0
-  - libsndfile=1.0.31=h9c3ff4c_1
-  - libsqlite=3.39.4=h753d276_0
-  - libssh2=1.10.0=haa6b8db_3
-  - libstdcxx-ng=12.1.0=ha89aaad_16
-  - libtiff=4.4.0=h0e0dad5_3
-  - libtool=2.4.6=h9c3ff4c_1008
-  - libudev1=249=h166bdaf_4
-  - libuuid=2.32.1=h7f98852_1000
-  - libvorbis=1.3.7=h9c3ff4c_0
-  - libwebp-base=1.2.4=h166bdaf_0
-  - libxcb=1.13=h7f98852_1004
-  - libxkbcommon=1.0.3=he3ba5ed_0
-  - libxml2=2.10.2=h4c7fe37_1
-  - libzlib=1.2.12=h166bdaf_4
-  - locket=1.0.0=pyhd8ed1ab_0
-  - lz4=4.0.0=py310h5d5e884_2
-  - lz4-c=1.9.3=h9c3ff4c_1
-  - markdown=3.4.1=pyhd8ed1ab_0
-  - markupsafe=2.1.1=py310h5764c6d_1
-  - matplotlib=3.6.0=py310hff52083_0
-  - matplotlib-base=3.6.0=py310h8d5ebf3_0
-  - msgpack-python=1.0.4=py310hbf28c38_0
-  - multidict=6.0.2=py310h5764c6d_1
-  - munkres=1.1.4=pyh9f0ad1d_0
-  - mysql-common=8.0.30=haf5c9bc_1
-  - mysql-libs=8.0.30=h28c427c_1
-  - ncurses=6.3=h27087fc_1
-  - networkx=2.8.7=pyhd8ed1ab_0
-  - nomkl=1.0=h5ca1d4c_0
-  - nspr=4.32=h9c3ff4c_1
-  - nss=3.78=h2350873_0
-  - numcodecs=0.10.2=py310hd8f1fbe_0
-  - numexpr=2.8.3=py310hf05e7a9_100
-  - numpy=1.23.3=py310h53a5b5f_0
-  - oauthlib=3.2.1=pyhd8ed1ab_0
-  - openjpeg=2.5.0=h7d73246_1
-  - openssl=1.1.1q=h166bdaf_0
-  - opt_einsum=3.3.0=pyhd8ed1ab_1
-  - packaging=21.3=pyhd8ed1ab_0
-  - pandas=1.5.0=py310h769672d_0
-  - partd=1.3.0=pyhd8ed1ab_0
-  - patsy=0.5.2=pyhd8ed1ab_0
-  - pcre2=10.37=hc3806b6_1
-  - perl=5.32.1=2_h7f98852_perl5
-  - pillow=9.2.0=py310hbd86126_2
-  - pip=22.2.2=pyhd8ed1ab_0
-  - ply=3.11=py_1
-  - pomegranate=0.14.8=py310hc4a4660_0
-  - portaudio=19.6.0=h8e90077_6
-  - protobuf=4.21.7=py310hd8f1fbe_0
-  - psutil=5.9.2=py310h5764c6d_0
-  - pthread-stubs=0.4=h36c2ea0_1001
-  - pulseaudio=14.0=h0868958_9
-  - pyasn1=0.4.8=py_0
-  - pyasn1-modules=0.2.7=py_0
-  - pycparser=2.21=pyhd8ed1ab_0
-  - pyjwt=2.5.0=pyhd8ed1ab_0
-  - pyopenssl=22.0.0=pyhd8ed1ab_1
-  - pyparsing=3.0.9=pyhd8ed1ab_0
-  - pyqt=5.15.7=py310h29803b5_0
-  - pyqt5-sip=12.11.0=py310hd8f1fbe_0
-  - pysocks=1.7.1=pyha2e5f31_6
-  - python=3.10.6=h582c2e5_0_cpython
-  - python-dateutil=2.8.2=pyhd8ed1ab_0
-  - python-flatbuffers=2.0=pyhd8ed1ab_0
-  - python_abi=3.10=2_cp310
-  - pytz=2022.4=pyhd8ed1ab_0
-  - pyu2f=0.1.5=pyhd8ed1ab_0
-  - pyyaml=6.0=py310h5764c6d_4
-  - qt-main=5.15.6=hc525480_0
-  - re2=2022.06.01=h27087fc_0
-  - readline=8.1.2=h0f457ee_0
-  - requests=2.28.1=pyhd8ed1ab_1
-  - requests-oauthlib=1.3.1=pyhd8ed1ab_0
-  - rsa=4.9=pyhd8ed1ab_0
-  - samtools=1.16.1=h6899075_0
-  - scikit-allel=1.3.5=py310hb5077e9_1
-  - scikit-learn=1.1.2=py310h0c3af53_0
-  - scipy=1.9.1=py310hdfbd76f_0
-  - seaborn=0.12.0=hd8ed1ab_0
-  - seaborn-base=0.12.0=pyhd8ed1ab_0
-  - setuptools=65.4.1=pyhd8ed1ab_0
-  - sip=6.6.2=py310hd8f1fbe_0
-  - six=1.16.0=pyh6c4a22f_0
-  - slim=3.7
-  - snappy=1.1.9=hbd366e4_1
-  - sortedcontainers=2.4.0=pyhd8ed1ab_0
-  - sqlite=3.39.4=h4ff8645_0
-  - statsmodels=0.13.2=py310hde88566_0
-  - tblib=1.7.0=pyhd8ed1ab_0
-  - tensorboard=2.10.1=pyhd8ed1ab_0
-  - tensorboard-data-server=0.6.0=py310h597c629_2
-  - tensorboard-plugin-wit=1.8.1=pyhd8ed1ab_0
-  - tensorflow=2.10.0=cpu_py310hd1aba9c_0
-  - tensorflow-base=2.10.0=cpu_py310hc537a0e_0
-  - tensorflow-estimator=2.10.0=cpu_py310had6d012_0
-  - termcolor=2.0.1=pyhd8ed1ab_1
-  - threadpoolctl=3.1.0=pyh8a188c0_0
-  - tk=8.6.12=h27826a3_0
-  - toml=0.10.2=pyhd8ed1ab_0
-  - toolz=0.12.0=pyhd8ed1ab_0
-  - tornado=6.1=py310h5764c6d_3
-  - tqdm=4.64.1=pyhd8ed1ab_0
-  - typing-extensions=4.3.0=hd8ed1ab_0
-  - typing_extensions=4.3.0=pyha770c72_0
-  - tzdata=2022d=h191b570_0
-  - unicodedata2=14.0.0=py310h5764c6d_1
-  - urllib3=1.26.11=pyhd8ed1ab_0
-  - werkzeug=2.2.2=pyhd8ed1ab_0
-  - wheel=0.37.1=pyhd8ed1ab_0
-  - wrapt=1.14.1=py310h5764c6d_0
-  - xcb-util=0.4.0=h516909a_0
-  - xcb-util-image=0.4.0=h166bdaf_0
-  - xcb-util-keysyms=0.4.0=h516909a_0
-  - xcb-util-renderutil=0.3.9=h166bdaf_0
-  - xcb-util-wm=0.4.1=h516909a_0
-  - xorg-libxau=1.0.9=h7f98852_0
-  - xorg-libxdmcp=1.1.3=h7f98852_0
-  - xz=5.2.6=h166bdaf_0
-  - yaml=0.2.5=h7f98852_2
-  - yarl=1.7.2=py310h5764c6d_2
-  - zarr=2.13.2=pyhd8ed1ab_1
-  - zict=2.2.0=pyhd8ed1ab_0
-  - zipp=3.8.1=pyhd8ed1ab_0
-  - zlib=1.2.12=h166bdaf_4
-  - zstd=1.5.2=h6239696_4
diff --git a/build/lib/timesweeper/__init__.py b/build/lib/timesweeper/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/build/lib/timesweeper/__main__.py b/build/lib/timesweeper/__main__.py
new file mode 100644
index 0000000..e69de29
diff --git a/build/lib/timesweeper/cli.py b/build/lib/timesweeper/cli.py
new file mode 100644
index 0000000..a0035d2
--- /dev/null
+++ b/build/lib/timesweeper/cli.py
@@ -0,0 +1,480 @@
+import argparse
+import multiprocessing as mp
+import os
+import sys
+
+
+def ts_main():
+    agp = argparse.ArgumentParser(description="Timesweeper CLI")
+    subparsers = agp.add_subparsers(help="Timesweeper", dest="mode")
+
+    # simulate_stdpopsim.py
+    sim_s_parser = subparsers.add_parser(
+        name="sim_stdpopsim",
+        help="Injects time-series sampling into stdpopsim SLiM output script.",
+    )
+    sim_s_parser.add_argument(
+        "-v",
+        "--verbose",
+        required=False,
+        action="store_true",
+        dest="verbose",
+        help="Print verbose logging during process.",
+    )
+    sim_s_parser.add_argument(
+        "--threads",
+        required=False,
+        type=int,
+        default=mp.cpu_count() - 1,
+        dest="threads",
+        help="Number of processes to parallelize across.",
+    )
+    sim_s_parser.add_argument(
+        "--rep-range",
+        required=False,
+        dest="rep_range",
+        nargs=2,
+        help="<start, stop>. If used, only range(start, stop) will be simulated for reps. \
+            This is to allow for easy SLURM parallel simulations.",
+    )
+    sim_s_parser.add_argument(
+        metavar="YAML CONFIG",
+        dest="yaml_file",
+        help="YAML config file with all options defined.",
+    )
+
+    # simulate_custom.py
+    sim_c_parser = subparsers.add_parser(
+        name="sim_custom",
+        help="Simulates selection for training Timesweeper using a pre-made SLiM script.",
+    )
+    sim_c_parser.add_argument(
+        "--threads",
+        required=False,
+        type=int,
+        default=mp.cpu_count(),
+        dest="threads",
+        help="Number of processes to parallelize across. Defaults to all.",
+    )
+    sim_c_parser.add_argument(
+        "--rep-range",
+        required=False,
+        dest="rep_range",
+        nargs=2,
+        help="<start, stop>. If used, only range(start, stop) will be simulated for reps. \
+            This is to allow for easy SLURM parallel simulations.",
+    )
+    sim_c_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # make_training_features.py
+    mtf_parser = subparsers.add_parser(
+        name="condense",
+        help="Creates training data from simulated merged vcfs after process_vcfs.py has been run.",
+    )
+    mtf_parser = subparsers.add_parser(
+        name="condense",
+        help="Creates training data from simulated merged vcfs after process_vcfs.py has been run.",
+    )
+    mtf_parser.add_argument(
+        "--threads",
+        required=False,
+        type=int,
+        default=mp.cpu_count() - 1,
+        dest="threads",
+        help="Number of processes to parallelize across.",
+    )
+    mtf_parser.add_argument(
+        "-o",
+        "--outfile",
+        required=False,
+        type=str,
+        default="training_data.pkl",
+        dest="outfile",
+        help="Pickle file to dump dictionaries with training data to. Should probably end with .pkl.",
+    )
+    mtf_parser.add_argument(
+        "--subsample-inds",
+        required=False,
+        type=int,
+        dest="subsample_inds",
+        help="Number of individuals to subsample if using a larger simulation than needed for efficiency. NOTE: If you use this, keep the sample_sizes entry in the yaml config identical to the original simulations. This is mostly for the paper experiments and as such isn't very cleanly implemented.",
+    )
+    mtf_parser.add_argument(
+        "--subsample-tps",
+        required=False,
+        type=int,
+        dest="subsample_tps",
+        help="Number of timepoints to subsample if using a larger simulation than needed for efficiency. NOTE: If you use this, keep the sample_sizes entry in the yaml config identical to the original simulations. This is mostly for the paper experiments and as such isn't very cleanly implemented.",
+    )
+    mtf_parser.add_argument(
+        "--og-tps",
+        required=False,
+        type=int,
+        dest="og_tps",
+        help="Number of timepoints taken in the original data to be subsetted if using --subample-tps.",
+    )
+    mtf_parser.add_argument(
+        "-m",
+        "--missingness",
+        metavar="MISSINGNESS",
+        dest="missingness",
+        type=float,
+        required=False,
+        default=0.0,
+        help="Missingness rate in range of [0,1], used as the parameter of a binomial distribution for randomly removing known values.",
+    )
+    mtf_parser.add_argument(
+        "-f",
+        "--freq-increase-threshold",
+        metavar="FREQ_INC_THRESHOLD",
+        dest="freq_inc_thr",
+        type=float,
+        required=False,
+        default=0.0,
+        help="If given, only include sim replicates where the sweep site has a minimum increase of <freq_inc_thr> from the first timepoint to the last.",
+    )
+    mtf_parser.add_argument(
+        "-a",
+        "--allow-shoulders",
+        metavar="SAMPLING_OFFSET",
+        dest="allow_shoulders",
+        type=int,
+        required=False,
+        help="If value given: sample a uniform distribution with bounds of given value to offset training sample from central SNP.",
+    )
+    mtf_parser.add_argument(
+        "--hft",
+        required=False,
+        action="store_true",
+        dest="hft",
+        help="Whether to calculate HFT alongside AFT. Computationally more expensive.",
+    )
+    mtf_parser.add_argument(
+        "--no-progress",
+        action="store_true",
+        dest="no_progress",
+        help="Turn off progress bar.",
+    )
+    mtf_parser.add_argument(
+        "--verbose",
+        action="store_true",
+        dest="verbose",
+        help="Raise warnings from issues usually stemming from bad replicates.",
+    )
+    mtf_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # train_nets.py
+    nets_parser = subparsers.add_parser(
+        name="train",
+        help="Handler script for neural network training and prediction for TimeSweeper Package.\
+            Will train two models: one for the series of timepoints generated using the hfs vectors over a timepoint and one ",
+    )
+    nets_parser.add_argument(
+        "-i",
+        "--training-data",
+        metavar="TRAINING_DATA",
+        dest="training_data",
+        type=str,
+        required=True,
+        help="Pickle file containing data formatted with make_training_features.py.",
+    )
+    nets_parser.add_argument(
+        "-d",
+        "--data-type",
+        required=True,
+        metavar="DATA_TYPE",
+        dest="data_type",
+        help="AFT or HFT data preparation.",
+    )
+    nets_parser.add_argument(
+        "-s",
+        "--subsample-amount",
+        metavar="SUBSAMPLE_AMOUNT",
+        dest="subsample_amount",
+        type=int,
+        required=False,
+        help="Amount of data to subsample for each class to test for sample size effects.",
+    )
+    nets_parser.add_argument(
+        "-n",
+        "--experiment-name",
+        metavar="EXPERIMENT_NAME",
+        dest="experiment_name",
+        type=str,
+        required=False,
+        default="ts_experiment",
+        help="Identifier for the experiment used to generate the data. Optional, but helpful in differentiating runs.",
+    )
+    nets_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+    nets_parser.add_argument(
+        "--single-tp",
+        dest="single_tp",
+        required=False,
+        action="store_true",
+        help="Whether to use the tp1_model module for a special case experiment.",
+    )
+
+    # find_sweeps.py
+    sweeps_parser = subparsers.add_parser(
+        name="detect",
+        help="Module for iterating across windows in a time-series vcf file and predicting whether a sweep is present at each snp-centralized window.",
+    )
+    sweeps_parser.add_argument(
+        "-i",
+        "--input-vcf",
+        dest="input_vcf",
+        help="Merged VCF to scan for sweeps. Must be merged VCF where files are merged in order from earliest to latest sampling time, -0 flag must be used.",
+        required=True,
+    )
+    sweeps_parser.add_argument(
+        "--benchmark",
+        dest="benchmark",
+        action="store_true",
+        help="If testing on simulated data and would like to report the mutation \
+            type stored by SLiM during outputVCFSample as well as neutral predictions, use this flag. \
+            Otherwise the mutation type will not be looked for in the VCF entry nor reported with results.",
+        required=False,
+    )
+    sweeps_parser.add_argument(
+        "--aft-model",
+        dest="aft_model",
+        help="Path to Keras2-style saved model to load for aft prediction.",
+        required=True,
+    )
+    sweeps_parser.add_argument(
+        "--hft-model",
+        dest="hft_model",
+        help="Path to Keras2-style saved model to load for hft prediction.",
+        required=False,
+    )
+    sweeps_parser.add_argument(
+        "-o",
+        "--out-dir",
+        dest="outdir",
+        help="Directory to write output to.",
+        required=True,
+    )
+    sweeps_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # plot_training_data.py
+    input_plot_parser = subparsers.add_parser(
+        name="plot_training",
+        help="Plots central SNPs from simulations to visually inspect mean trends over replicates.",
+    )
+    input_plot_parser.add_argument(
+        "-i",
+        "--input-pickle",
+        dest="input_pickle",
+        metavar="INPUT PICKLE",
+        type=str,
+        required=True,
+        help="Pickle file containing dictionary of structure dict[sweep][rep]['aft'] created by make_training_features.py.",
+    )
+    input_plot_parser.add_argument(
+        "-n",
+        "--schema-name",
+        metavar="SCHEMA NAME",
+        dest="schema_name",
+        required=False,
+        default="simulation_center_means",
+        type=str,
+        help="Experiment label to use for output file naming.",
+    )
+    input_plot_parser.add_argument(
+        "-o",
+        "--output",
+        metavar="OUTPUT DIR",
+        dest="output_dir",
+        required=False,
+        default=".",
+        type=str,
+        help="Directory to write images to.",
+    )
+    input_plot_parser.add_argument(
+        "--save-example",
+        dest="save_example",
+        required=False,
+        action="store_true",
+        help="Will create a directory with example input matrices.",
+    )
+
+    freq_plot_parser = subparsers.add_parser(
+        name="plot_freqs",
+        help="Create a bedfile of major and minor allele frequency changes over time.",
+    )
+    freq_plot_parser.add_argument(
+        "-i",
+        "--input",
+        dest="input",
+        metavar="INPUT VCF FILE",
+        type=str,
+        required=True,
+        help="Merged time-series VCF file to pull SNPs and frequencies from.",
+    )
+    freq_plot_parser.add_argument(
+        "-o",
+        "--output",
+        metavar="OUTPUT FILE PREFIX",
+        dest="output",
+        required=False,
+        default=sys.stdout,
+        type=str,
+        help="""Bedgraph file prefix, two files will be written using this prefix + '{.major,.minor}.bedGraph.
+        The '.minor' file denotes the 4th column is the frequency change from last to first timepoints of the allele with the largest change over those epochs.
+        The '.major' file denotes the 4th column is the frequency change from last to first timepoints of 1-minor allele at each SNP, with the 'minor' allele being described above as the highest-velocity allele across timepoints in a given SNP.
+        Example: 'ts_output/d_simulans_experiment_1'.
+        """,
+    )
+    freq_plot_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # parse_sim_logs.py
+    summarize_parser = subparsers.add_parser(
+        name="summarize",
+        help="Creates a CSV of data parsed from slim log files.",
+    )
+    summarize_parser.add_argument(
+        "--threads",
+        required=False,
+        type=int,
+        default=mp.cpu_count(),
+        dest="threads",
+        help="Number of processes to parallelize across. Defaults to all.",
+    )
+    summarize_parser.add_argument(
+        "-n",
+        "--experiment-name",
+        metavar="EXPERIMENT_NAME",
+        dest="experiment_name",
+        type=str,
+        required=False,
+        default="ts_experiment",
+        help="Identifier for the experiment used to generate the data. Optional, but helpful in differentiating runs.",
+    )
+    summarize_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # merge_locs_acc.py
+    merge_parser = subparsers.add_parser(
+        name="merge_logs",
+        help="Merges the summary TSV from SLiM logs with test data predictions.",
+    )
+    merge_parser.add_argument(
+        "-s",
+        "--summary-tsv",
+        dest="summary_tsv",
+        help="TSV created using `parse_slim_logs.py`.",
+        required=True,
+    )
+    merge_parser.add_argument(
+        "-i",
+        "--input-test-file",
+        dest="test_file",
+        help="Files to merge ",
+        required=True,
+    )
+    merge_parser.add_argument(
+        "-o",
+        "--output-dir",
+        metavar="OUTPUT DIR",
+        dest="output_dir",
+        required=True,
+        type=str,
+        help="Directory to write merged results and plots to.",
+    )
+
+    ua = agp.parse_args()
+
+    # fmt: off
+    if ua.mode == "sim_stdpopsim":
+        from timesweeper import simulate_stdpopsim
+        simulate_stdpopsim.main(ua)
+
+    elif ua.mode == "sim_custom":
+        from timesweeper import simulate_custom
+        simulate_custom.main(ua)
+
+    elif ua.mode == "condense":
+        from timesweeper import make_training_features
+        make_training_features.main(ua)    
+
+    elif ua.mode == "train":
+        if ua.single_tp:
+            from timesweeper import tp1_model
+            tp1_model.main(ua)   
+        elif "shoulder" in str(ua.experiment_name).lower():
+            from timesweeper import train_shoulder_nets
+            train_shoulder_nets.main(ua)  
+        elif "shic" in str(ua.experiment_name).lower():
+            from timesweeper import train_nets_shic
+            train_nets_shic.main(ua)
+        else:
+            from timesweeper import train_nets
+            train_nets.main(ua)
+        
+    elif ua.mode == "detect":
+        from timesweeper import find_sweeps_vcf as find_sweeps_vcf
+        find_sweeps_vcf.main(ua)   
+
+    elif ua.mode == "plot_training":
+        from timesweeper.plotting import plot_training_data as plot_training
+        plot_training.main(ua)   
+
+    elif ua.mode == "plot_freqs":
+        from timesweeper.plotting import create_freq_track as cf
+        cf.main(ua)
+        
+    elif ua.mode == "summarize":
+        from timesweeper import parse_slim_logs as psl
+        psl.main(ua)
+    
+    elif ua.mode == "merge_logs":
+        from timesweeper import merge_logs_acc as mla
+        mla.main(ua)
+        
+    elif ua.mode == None:
+        agp.print_help()
+
+
+if __name__ == "__main__":
+    ts_main()
diff --git a/build/lib/timesweeper/find_sweeps_npz.py b/build/lib/timesweeper/find_sweeps_npz.py
new file mode 100644
index 0000000..6875f45
--- /dev/null
+++ b/build/lib/timesweeper/find_sweeps_npz.py
@@ -0,0 +1,138 @@
+import argparse as ap
+import logging
+import os
+from itertools import cycle
+
+import numpy as np
+import pandas as pd
+import yaml
+from tensorflow.keras.models import load_model
+from tqdm import tqdm
+
+from utils.gen_utils import write_preds
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("timesweeper")
+logger.setLevel("INFO")
+
+pd.options.display.float_format = "{:.2f}".format
+
+
+def load_npz(npz_path):
+    npz_obj = np.load(npz_path)
+    aft = npz_obj["aftIn"]  # (snps, timepoints, window)
+    locs = npz_obj["aftInPosition"]  # (snps, window_locs)
+
+    return aft, locs
+
+
+def parse_npz_name(npz_path):
+    splitpath = npz_path.split("_")
+    chrom = splitpath[2]
+    rep = splitpath[-1].split(".")[0]
+
+    return chrom, rep
+
+
+def run_aft_windows(ts_aft, locs, chrom, model):
+    """
+    Iterates through windows of MAF time-series matrix and predicts using NN.
+    Args:
+        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP. Contains mut only if benchmarking == True.
+        genos (allel.GenotypeArray): Genotypes of all samples.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
+        model (Keras.model): Keras model to use for prediction.
+    Returns:
+        dict: Prediction values in the form of dict[snps[center]]
+        np.arr: the central-most window, either based on mutation type or closest to half size of chrom.
+    """
+    left_edges = list(locs[:, 0])
+    right_edges = list(locs[:, -1])
+    centers = list(locs[:, 25])
+    class_probs, sel_pred = model.predict(ts_aft)
+
+    results_list = zip(
+        cycle([chrom]), centers, left_edges, right_edges, class_probs, sel_pred
+    )
+
+    return results_list
+
+
+def load_nn(model_path, summary=False):
+    """
+    Loads the trained Keras network.
+    Args:
+        model_path (str): Path to Keras model.
+        summary (bool, optional): Whether to print out model summary or not. Defaults to False.
+    Returns:
+        Keras.model: Trained Keras model to use for prediction.
+    """
+    model = load_model(model_path)
+    if summary:
+        print(model.summary())
+
+    return model
+
+
+def write_preds(results_list, outfile, benchmark):
+    """
+    Writes NN predictions to file.
+    Args:
+        results_dict (dict): SNP NN prediction scores and window edges.
+        outfile (str): File to write results to.
+    """
+    lab_dict = {0: "Neut", 1: "SSV"}
+    chrom, centers, left_edges, right_edges, probs = zip(*results_list)
+
+    neut_scores = [i[0] for i in probs]
+    ssv_scores = [i[1] for i in probs]
+    classes = [lab_dict[np.argmax(i)] for i in probs]
+
+    predictions = pd.DataFrame(
+        {
+            "Chrom": chrom,
+            "BP": centers,
+            "Class": classes,
+            "Neut_Score": neut_scores,
+            "SSV_Score": ssv_scores,
+            "Win_Start": left_edges,
+            "Win_End": right_edges,
+        }
+    )
+
+    # predictions = predictions[predictions["Neut_Score"] < 0.5]
+
+    predictions.sort_values(["Chrom", "BP"], inplace=True)
+
+    predictions.to_csv(outfile, header=True, index=False, sep="\t", float_format="%.3f")
+
+    bed_df = predictions[["Chrom", "Win_Start", "Win_End", "BP"]]
+    bed_df.to_csv(outfile.replace(".csv", ".bed"), header=False, index=False, sep="\t")
+
+
+def main(ua):
+    outdir, aft_model = (
+        ua.outdir,
+        load_nn(ua.aft_model),
+    )
+
+    try:
+        if not os.path.exists(outdir):
+            os.makedirs(outdir)
+    except:
+        # running in high-parallel sometimes it errors when trying to check/create simultaneously
+        pass
+
+    # Load in everything
+    logger.info(f"Loading data from {ua.input_file}")
+    chrom, rep = parse_npz_name(ua.input_file)
+    ts_aft, locs = load_npz(ua.input_file)
+
+    # aft
+    logger.info("Predicting with AFT")
+    aft_predictions = run_aft_windows(ts_aft, locs, chrom, aft_model)
+    write_preds(aft_predictions, f"{outdir}/aft_{chrom}_{rep}_preds.csv", ua.benchmark)
+    logger.info(f"Done, results written to {outdir}/aft_{chrom}_{rep}_preds.csv")
diff --git a/build/lib/timesweeper/find_sweeps_vcf.py b/build/lib/timesweeper/find_sweeps_vcf.py
new file mode 100644
index 0000000..f42cfe7
--- /dev/null
+++ b/build/lib/timesweeper/find_sweeps_vcf.py
@@ -0,0 +1,181 @@
+import math
+import os
+
+import numpy as np
+from tensorflow.keras.models import load_model
+from tqdm import tqdm
+
+from timesweepermake_training_features import prep_ts_aft, get_window_idxs
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import read_config, write_preds, get_logger
+from timesweeper.utils import hap_utils as hu
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logger = get_logger("find_sweeps")
+
+
+def run_aft_windows(snps, genos, samp_sizes, win_size, model):
+    """
+    Iterates through windows of MAF time-series matrix and predicts using NN.
+
+    Args:
+        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP. Contains mut only if benchmarking == True.
+        genos (allel.GenotypeArray): Genotypes of all samples.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
+        model (Keras.model): Keras model to use for prediction.
+
+    Returns:
+        dict: Prediction values in the form of dict[snps[center]]
+        np.arr: the central-most window, either based on mutation type or closest to half size of chrom.
+    """
+    ts_aft = prep_ts_aft(genos, samp_sizes)
+
+    # Iterate over SNP windows and predict
+    buffer = math.floor(win_size / 2)
+
+    centers = range(buffer, len(snps) - buffer)
+    left_edges = []
+    right_edges = []
+    data = []
+    for center in tqdm(centers, desc="Predicting on AFT windows"):
+        try:
+            win_idxs = get_window_idxs(center, win_size)
+            window = ts_aft[:, win_idxs]
+            data.append(window)
+            left_edges.append(snps[win_idxs[0]][1])
+            right_edges.append(snps[win_idxs[-1]][1])
+
+        except Exception as e:
+            logger.warning(f"Center {snps[center]} raised error {e}")
+
+    class_probs, sel_pred = model.predict(np.stack(data))
+
+    results_dict = {}
+    for center, class_probs, sel_pred, l_e, r_e in zip(
+        centers, class_probs, sel_pred, left_edges, right_edges
+    ):
+        results_dict[snps[center]] = (class_probs, sel_pred, l_e, r_e)
+
+    return results_dict
+
+
+def run_hft_windows(snps, haps, ploidy, samp_sizes, win_size, model):
+    """
+    Iterates through windows of MAF time-series matrix and predicts using NN.
+    Args:
+        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP. Contains mut only if benchmarking == True.
+        haps (np.arr): Haplotypes of all samples.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
+        model (Keras.model): Keras model to use for prediction.
+    Returns:
+        dict: Prediction values in the form of dict[snps[center]]
+        np.arr: the central-most window, either based on mutation type or closest to half size of chrom.
+    """
+    results_dict = {}
+    buffer = math.floor(win_size / 2)
+    centers = range(buffer, len(snps) - buffer)
+    left_edges = []
+    right_edges = []
+    data = []
+    for center in tqdm(centers, desc="Predicting on HFS windows"):
+        try:
+            win_idxs = get_window_idxs(center, win_size)
+            window = np.swapaxes(haps[win_idxs, :], 0, 1)
+            str_window = hu.haps_to_strlist(window)
+            hfs = hu.getTSHapFreqs(str_window, [i * ploidy for i in samp_sizes])
+            data.append(hfs)
+            left_edges.append(snps[win_idxs[0]][1])
+            right_edges.append(snps[win_idxs[-1]][1])
+
+        except Exception as e:
+            logger.warning(f"Center {snps[center]} raised error {e}")
+            continue
+
+    probs = model.predict(np.stack(data))
+
+    results_dict = {}
+    for center, prob, l_e, r_e in zip(centers, probs, left_edges, right_edges):
+        results_dict[snps[center]] = (prob, l_e, r_e)
+
+    return results_dict
+
+
+def get_window_idxs(center_idx, win_size):
+    """
+    Gets the win_size number of snps around a central snp.
+
+    Args:
+        center_idx (int): Index of the central SNP to use for the window.
+        win_size (int): Size of window to use around the SNP, optimally odd number.
+
+    Returns:
+        list: Indices of all SNPs to grab for the feature matrix.
+    """
+    half_window = math.floor(win_size / 2)
+    return list(range(center_idx - half_window, center_idx + half_window + 1))
+
+
+def load_nn(model_path, summary=False):
+    """
+    Loads the trained Keras network.
+
+    Args:
+        model_path (str): Path to Keras model.
+        summary (bool, optional): Whether to print out model summary or not. Defaults to False.
+
+    Returns:
+        Keras.model: Trained Keras model to use for prediction.
+    """
+    model = load_model(model_path)
+    if summary:
+        print(model.summary())
+
+    return model
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    (work_dir, samp_sizes, ploidy, win_size, outdir, aft_model) = (
+        yaml_data["work dir"],
+        yaml_data["sample sizes"],
+        yaml_data["ploidy"],
+        yaml_data["win_size"],
+        ua.outdir,
+        load_nn(ua.aft_model),
+    )
+
+    outdir = ua.outdir
+    if not os.path.exists(outdir):
+        os.makedirs(outdir)
+
+    # Chunk and iterate for NN predictions to not take up too much space
+    vcf_iter = su.get_vcf_iter(ua.input_vcf, ua.benchmark)
+    for chunk_idx, chunk in enumerate(vcf_iter):
+        chunk = chunk[0]  # Why you gotta do me like that, skallel?
+        logger.info(f"Processing VCF chunk {chunk_idx}")
+
+        # aft
+        try:
+            genos, snps = su.vcf_to_genos(chunk, ua.benchmark)
+            aft_predictions = run_aft_windows(
+                snps, genos, samp_sizes, win_size, aft_model
+            )
+            write_preds(aft_predictions, f"{outdir}/aft_preds.csv", ua.benchmark)
+
+        except Exception as e:
+            logger.error(f"Cannot process chunk {chunk_idx} using AFT due to {e}")
+
+        # hft
+        if ua.hft_model:
+            try:
+                haps, snps = su.vcf_to_haps(chunk, ua.benchmark)
+                aft_predictions = run_hft_windows(
+                    snps, haps, ploidy, samp_sizes, win_size, load_nn(ua.hft_model)
+                )
+                write_preds(aft_predictions, f"{outdir}/hft_preds.csv", ua.benchmark)
+
+            except Exception as e:
+                logger.error(f"Cannot process chunk {chunk_idx} using HFT due to {e}")
diff --git a/build/lib/timesweeper/make_training_features.py b/build/lib/timesweeper/make_training_features.py
new file mode 100644
index 0000000..7bce9ec
--- /dev/null
+++ b/build/lib/timesweeper/make_training_features.py
@@ -0,0 +1,450 @@
+import logging
+import math
+import multiprocessing as mp
+import os
+import pickle
+import sys
+from glob import glob
+from itertools import cycle
+from random import sample
+import allel
+import numpy as np
+from numpy.random import default_rng
+from tqdm import tqdm
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import get_rep_id, get_scenario_from_filename, read_config
+from timesweeper.utils.hap_utils import getTSHapFreqs, haps_to_strlist
+
+logging.basicConfig()
+logger = logging.getLogger("make_training_feats")
+logger.setLevel("INFO")
+
+import warnings
+
+warnings.filterwarnings("error")
+
+
+def draw_rand_center_offset(max_offset=7000):
+    rng = default_rng(np.random.seed(int.from_bytes(os.urandom(4), byteorder="little")))
+    return int(rng.uniform(-max_offset, max_offset, 1)[0])
+
+
+def add_missingness(data, m_rate, nan_val=-1):
+    rng = default_rng(np.random.seed(int.from_bytes(os.urandom(4), byteorder="little")))
+    missing = rng.binomial(1, m_rate, data.shape)
+    data[missing == 1] = nan_val
+
+    return data
+
+
+def check_freq_increase(ts_afs, min_increase=0.25):
+    """Quick test to make sure a given example is increasing properly, can be used to filter training."""
+    center = int(ts_afs.shape[0] / 2)
+    if ts_afs[-1, center] - ts_afs[0, center] >= min_increase:
+        return True
+    else:
+        return False
+
+
+def subsample_tps(inds_per_tp, og_tps, num_sampled_tps):
+    """Subsamples timepoints during VCF by sampling samples list."""
+    samples_list = []
+    if num_sampled_tps == 1:
+        t = og_tps + 1  # VCFs merged samples are 1-indexed, get last
+        samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+    elif num_sampled_tps == 2:
+        # This covers the first sampling point
+        samples_list.extend([f"i{i}" for i in range(inds_per_tp)])
+        t = og_tps + 1  # VCFs merged samples are 1-indexed, get last
+        samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+    else:
+        tps = np.linspace(1, og_tps, num_sampled_tps, dtype=int)
+        # This covers the first sampling point
+        samples_list.extend([f"i{i}" for i in range(inds_per_tp)])
+        # This gets the rest
+        for t in tps[1:]:
+            samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+
+    return samples_list
+
+
+def subsample_inds(inds_per_tp, subsample_size, num_tps):
+    """Subsamples individuals from each timepoint to re-use simulations."""
+    samples_list = []
+    samples_list.extend(
+        [f"i{i}" for i in np.random.choice(inds_per_tp, subsample_size, replace=False)]
+    )
+    for s in range(2, num_tps + 1):
+        samples_list.extend(
+            [
+                f"{s}:i{i}"
+                for i in np.random.choice(inds_per_tp, subsample_size, replace=False)
+            ]
+        )
+    return samples_list
+
+
+def prep_ts_aft(genos, samp_sizes):
+    """
+    Iterates through timepoints and creates MAF feature matrices.
+
+    Args:
+        genos (allel.GenotypeArray): Genotype array containing all timepoints.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+
+    Returns:
+        np.arr: MAF array to use for predictions. Shape is (timepoints, MAF).
+    """
+    # Prep genos into time-series format and calculate Maft
+    ts_genos = su.split_arr(genos, samp_sizes)
+
+    min_alleles, first_genos, last_genos = su.get_vel_minor_alleles(
+        ts_genos, np.max(genos)
+    )
+
+    # mid_idx = np.floor(len(first_genos) / 2)
+
+    ts_maft = []
+    for timepoint in ts_genos:
+        _genos = []
+        _genotypes = allel.GenotypeArray(timepoint).count_alleles(
+            max_allele=min_alleles.max()
+        )
+
+        for snp, min_allele_idx in zip(_genotypes, min_alleles):
+            maf = su.calc_maft(snp, min_allele_idx)
+            _genos.append(maf)
+
+        ts_maft.append(_genos)
+
+    return np.stack(ts_maft)
+
+
+def get_window_idxs(center_idx, win_size):
+    """
+    Gets the win_size number of snps around a central snp.
+
+    Args:
+        center_idx (int): Index of the central SNP to use for the window.
+        win_size (int): Size of window to use around the SNP, optimally odd number.
+
+    Returns:
+        list: Indices of all SNPs to grab for the feature matrix.
+    """
+    half_window = math.floor(win_size / 2)
+    return list(range(center_idx - half_window, center_idx + half_window + 1))
+
+
+def get_aft_central_window(
+    snps, genos, samp_sizes, win_size, missingness, mut_types, offset
+):
+    """
+    Check for whether a non-control mutation type is present. If not, return the central-most mutation.
+
+    Args:
+        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP.
+        genos (allel.GenotypeArray): Genotypes of all samples.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
+        scenario (str): Entry from the scenarios config option.
+        missingness (float): Parameter of binomial distribution to pull missingness from.
+        mut_types (list[int]): List of mutation types that are not considered the "control" case.
+    Returns:
+        np.arr: The central-most window, either based on mutation type or closest to half size of chrom.
+        float: Selection coefficient.
+    """
+    ts_aft = prep_ts_aft(genos, samp_sizes)
+
+    buffer = int(win_size / 2)
+    centers = range(buffer, len(snps) - buffer)
+
+    center_idx = int(len(centers) / 2)
+
+    try:
+        for center in centers:
+            if snps[center][2] in mut_types:
+                center_idx = center
+                break
+            else:
+                pass
+    except:
+        pass
+
+    sel_coeff = snps[center_idx][3]
+
+    # If a sweep and offset lands inside window still use center. Otherwise shoulder or neut.
+    if offset > 0:
+        rand_offset = draw_rand_center_offset()
+        if (snps[center_idx][2] in mut_types) and (abs(rand_offset) > win_size):
+            rand_offset = 0
+        else:
+            center_idx += rand_offset
+    else:
+        rand_offset = 0
+
+    win_idxs = get_window_idxs(center_idx, win_size)
+
+    window = ts_aft[:, win_idxs]
+    center_aft = window
+
+    missing_center_aft = add_missingness(
+        center_aft, m_rate=missingness
+    )  # If no missingness, will just return
+
+    return missing_center_aft, sel_coeff, rand_offset
+
+
+def get_hft_central_window(snps, haps, samp_sizes, win_size, mut_types, offset):
+    """
+    Iterates through windows of MAF time-series matrix and gets the central window.
+    Does not have as many utility functions as AFT such as missingness and variable sorting methods.
+
+    Args:
+        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP.
+        haps (np.arr): Haplotypes of all samples.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
+    Returns:
+        np.arr: The central-most window, either based on mutation type or closest to half size of chrom.
+    """
+    buffer = int(win_size / 2)
+    centers = range(buffer, len(snps) - buffer)
+
+    center_idx = int(len(centers) / 2)
+
+    try:
+        for center in centers:
+            if snps[center][2] in mut_types:
+                center_idx = center
+                break
+            else:
+                pass
+    except:
+        pass
+
+    sel_coeff = snps[center_idx][3]
+
+    # If a sweep and offset lands inside window still use center. Otherwise shoulder or neut.
+    if offset > 0:
+        rand_offset = draw_rand_center_offset()
+        if (snps[center_idx][2] in mut_types) and (abs(rand_offset) > win_size):
+            rand_offset = 0
+        else:
+            center_idx += rand_offset
+    else:
+        rand_offset = 0
+
+    win_idxs = get_window_idxs(center_idx, win_size)
+    window = np.swapaxes(haps[win_idxs, :], 0, 1)
+    str_window = haps_to_strlist(window)
+    central_hfs = getTSHapFreqs(str_window, samp_sizes)
+
+    return central_hfs, sel_coeff, rand_offset
+
+
+def aft_worker(
+    in_vcf,
+    mut_types,
+    scenarios,
+    samp_sizes,
+    samps_list,
+    win_size,
+    offset,
+    missingness,
+    verbose=False,
+):
+    benchmark = True  # Want to get all the info we can from sims in training
+    try:
+        id = get_rep_id(in_vcf)
+        scenario = get_scenario_from_filename(in_vcf, scenarios)
+
+        vcf = su.read_vcf(in_vcf, samps_list, benchmark)
+        genos, snps = su.vcf_to_genos(vcf, benchmark)
+
+        central_aft, sel_coeff, rand_offset = get_aft_central_window(
+            snps, genos, samp_sizes, win_size, missingness, mut_types, offset
+        )
+
+        return id, scenario, central_aft, sel_coeff, rand_offset
+
+    except UserWarning as Ue:
+        print(Ue)
+        return None
+
+    except Exception as e:
+        if verbose:
+            logger.warning(f"Could not process {in_vcf}")
+            logger.warning(f"Exception: {e}")
+            sys.stdout.flush()
+            sys.stderr.flush()
+        return None
+
+
+def hft_worker(
+    in_vcf,
+    mut_types,
+    scenarios,
+    samp_sizes,
+    samps_list,
+    win_size,
+    offset,
+    ploidy=1,
+    verbose=False,
+):
+    benchmark = True
+    try:
+        id = get_rep_id(in_vcf)
+        scenario = get_scenario_from_filename(in_vcf, scenarios)
+
+        vcf = su.read_vcf(in_vcf, samps_list, benchmark)
+        haps, snps = su.vcf_to_haps(vcf, benchmark)
+
+        central_hft, sel_coeff, rand_offset = get_hft_central_window(
+            snps, haps, [ploidy * i for i in samp_sizes], win_size, mut_types, offset
+        )
+
+        return id, scenario, central_hft, sel_coeff, rand_offset
+
+    except UserWarning as Ue:
+        # print(Ue)
+        return None
+
+    except Exception as e:
+        if verbose:
+            logger.warning(f"Could not process {in_vcf}")
+            logger.warning(f"Exception: {e}")
+            sys.stdout.flush()
+            sys.stderr.flush()
+        return None
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    scenarios, mut_types, work_dir, samp_sizes, ploidy, win_size, threads = (
+        yaml_data["scenarios"],
+        yaml_data["mut types"],
+        yaml_data["work dir"],
+        yaml_data["sample sizes"],
+        yaml_data["ploidy"],
+        yaml_data["win_size"],
+        ua.threads,
+    )
+    if ua.subsample_inds:
+        if ua.subsample_tps:
+            logger.error("Can't subsample both timepoints and individuals.")
+
+        samps_list = subsample_inds(samp_sizes[0], ua.subsample_inds, len(samp_sizes))
+    else:
+        samps_list = None
+
+    if ua.subsample_tps:
+        if not ua.og_tps:
+            logger.error(
+                "Must provide original simulation number of tps in order to properly subset."
+            )
+
+        samps_list = subsample_tps(samp_sizes[0], ua.og_tps, ua.subsample_tps)
+
+    if ua.allow_shoulders:
+        offset = int(ua.allow_shoulders)
+    else:
+        offset = 0
+
+    filelist = glob(f"{work_dir}/vcfs/*/*/merged.vcf", recursive=True)
+
+    aft_work_args = zip(
+        filelist,
+        cycle([mut_types]),
+        cycle([scenarios]),
+        cycle([samp_sizes]),
+        cycle([samps_list]),
+        cycle([win_size]),
+        cycle([offset]),
+        cycle([ua.missingness]),
+        cycle([ua.verbose]),
+    )
+    hft_work_args = zip(
+        filelist,
+        cycle([mut_types]),
+        cycle([scenarios]),
+        cycle([samp_sizes]),
+        cycle([samps_list]),
+        cycle([win_size]),
+        cycle([offset]),
+        cycle([int(ploidy)]),
+        cycle([ua.verbose]),
+    )
+
+    debug = False
+    if debug:
+        aft_work_res = [aft_worker(*i) for i in aft_work_args]
+        hft_work_res = [hft_worker(*i) for i in hft_work_args]
+    else:
+        pool = mp.Pool(threads)
+        if ua.no_progress:
+            aft_work_res = pool.starmap(
+                aft_worker,
+                aft_work_args,
+                chunksize=4,
+            )
+
+            if ua.hft:
+                hft_work_res = pool.starmap(
+                    hft_worker,
+                    hft_work_args,
+                    chunksize=4,
+                )
+
+            pool.close()
+        else:
+            aft_work_res = pool.starmap(
+                aft_worker,
+                tqdm(
+                    aft_work_args,
+                    desc="Formatting AFT training data",
+                    total=len(filelist),
+                ),
+                chunksize=4,
+            )
+            if ua.hft:
+                hft_work_res = pool.starmap(
+                    hft_worker,
+                    tqdm(
+                        hft_work_args,
+                        desc="Formatting HFT training data",
+                        total=len(filelist),
+                    ),
+                    chunksize=4,
+                )
+            pool.close()
+
+    pickle_dict = {}
+    for s in scenarios:
+        pickle_dict[s] = {}
+
+    for res in aft_work_res:
+        if res:
+            rep, scenario, aft, s, off = res
+
+            pickle_dict[scenario][rep] = {}
+            pickle_dict[scenario][rep]["aft"] = aft
+            pickle_dict[scenario][rep]["sel_coeff"] = s
+            pickle_dict[scenario][rep]["center_offset"] = off
+
+    if ua.hft:
+        for res in hft_work_res:
+            try:
+                if res:
+                    rep, scenario, hft, s, off = res
+                    if rep in pickle_dict[scenario].keys():
+                        pickle_dict[scenario][rep]["hft"] = hft
+                        pickle_dict[scenario][rep]["sel_coeff"] = s
+                        pickle_dict[scenario][rep]["center_offset"] = off
+
+            except KeyError as e:
+                print(e)
+                print(res)
+                pass
+
+    with open(ua.outfile, "wb") as outfile:
+        pickle.dump(pickle_dict, outfile)
diff --git a/build/lib/timesweeper/merge_logs_acc.py b/build/lib/timesweeper/merge_logs_acc.py
new file mode 100644
index 0000000..b5ff396
--- /dev/null
+++ b/build/lib/timesweeper/merge_logs_acc.py
@@ -0,0 +1,66 @@
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+
+def main(ua):
+    if not os.path.exists(ua.output_dir):
+        os.makedirs(ua.output_dir)
+
+    log_df = pd.read_csv(ua.summary_tsv, header=0, sep="\t")
+    acc_df = pd.read_csv(ua.test_file, header=0, sep=",")
+    merged = pd.merge(
+        right=acc_df,
+        left=log_df,
+        how="right",
+        left_on=["rep", "sweep"],
+        right_on=["rep", "class"],
+    )
+    merged["abs_error"] = np.abs(merged["pred_sel_coeff"] - merged["true_sel_coeff"])
+    merged.to_csv(
+        f"{ua.output_dir}/merged_table.tsv", index=False, header=True, sep="\t"
+    )
+
+    merged.plot(
+        x="selCoeff",
+        y="abs_error",
+        kind="scatter",
+        title="Absolute S Prediction Error vs Selection Coefficient",
+        xlabel="Selection Coefficient s",
+        ylabel="Absolute Error | abs(predicted s - true s)",
+        xlim=(0.0, 0.25),
+        ylim=(0.0, 0.25),
+    )
+    plt.savefig(f"{ua.output_dir}/error_vs_selcoeff.png")
+    plt.clf()
+
+    merged.plot(
+        x="sampOffset",
+        y="abs_error",
+        kind="scatter",
+        title="Absolute S Prediction Error vs Selection Coefficient",
+        xlabel="Sampling Offset from Selection",
+        ylabel="Absolute Error | abs(predicted s - true s)",
+        xlim=(-50, 50),
+        ylim=(0.0, 0.25),
+        xticks=range(-50, 60, 10),
+    )
+    plt.savefig(f"{ua.output_dir}/error_vs_sampletime.png")
+    plt.clf()
+
+    merged.plot(
+        x="sampOffset",
+        y="selCoeff",
+        kind="scatter",
+        title="Sampling Offset from Selection vs Selection Coefficient",
+        xlabel="Sampling Offset from Selection",
+        ylabel="Selection Coefficient",
+        xlim=(-50, 50),
+        ylim=(0.0, 0.25),
+        xticks=range(-50, 60, 10),
+    )
+    plt.savefig(f"{ua.output_dir}/selcoeff_vs_sampletime.png")
+    plt.clf()
diff --git a/build/lib/timesweeper/models.py b/build/lib/timesweeper/models.py
new file mode 100644
index 0000000..77c1254
--- /dev/null
+++ b/build/lib/timesweeper/models.py
@@ -0,0 +1,213 @@
+from tensorflow.keras import layers
+from tensorflow.keras.models import Model
+
+# fmt: off
+def create_TS_class_model(datadim, n_class):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Conv1D(64, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv1D(64, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling1D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(512, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="Timesweeper_Class")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+def create_TS_reg_model(datadim):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Conv1D(64, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv1D(64, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling1D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    reg_output = layers.Dense(1, activation="sigmoid", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="Timesweeper_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
+
+# fmt: off
+def create_1tp_class_model(datadim, n_class):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+
+    h = layers.Dense(264, activation="relu")(model_in)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(56, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="Timesweeper_Class")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+
+def create_1tp_reg_model(datadim):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Dense(264, activation="relu")(model_in)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(56, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+
+    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="Timesweeper_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
+# fmt: on
+
+def create_rnn_reg_model(datadim):
+    model_in = layers.Input(datadim)
+    h = layers.LSTM(100, activation="relu")(model_in)
+    h = layers.LSTM(0.2)(h)        
+    h = layers.LSTM(100, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+
+    reg_output = layers.Dense(1, activation="sigmoid", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="RNN_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
+
+
+def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
+    # Attention and Normalization
+    x = layers.MultiHeadAttention(
+        key_dim=head_size, num_heads=num_heads, dropout=dropout
+    )(inputs, inputs)
+    x = layers.Dropout(dropout)(x)
+    x = layers.LayerNormalization(epsilon=1e-6)(x)
+    res = x + inputs
+
+    # Feed Forward Part
+    x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(res)
+    x = layers.Dropout(dropout)(x)
+    x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
+    x = layers.LayerNormalization(epsilon=1e-6)(x)
+
+    return x + res
+
+
+def create_transformer_reg_model(
+    input_shape,
+    head_size,
+    num_heads,
+    ff_dim,
+    num_transformer_blocks,
+    mlp_units,
+    dropout=0,
+    mlp_dropout=0,
+):
+    inputs = layers.Input(shape=input_shape)
+    x = inputs
+    for _ in range(num_transformer_blocks):
+        x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
+
+    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
+    for dim in mlp_units:
+        x = layers.Dense(dim, activation="relu")(x)
+        x = layers.Dropout(mlp_dropout)(x)
+    outputs = layers.Dense(1, activation="relu")(x)
+
+    model = Model(inputs, outputs, name="Timesweeper_Transformer_Reg")
+
+    model.compile(
+        loss="mse",
+        optimizer="adam",
+        metrics="mse",
+    )
+
+    return model
+
+
+def create_transformer_class_model(
+    input_shape,
+    head_size,
+    num_heads,
+    ff_dim,
+    num_transformer_blocks,
+    mlp_units,
+    dropout=0,
+    mlp_dropout=0,
+    n_class=3,
+):
+    inputs = layers.Input(shape=input_shape)
+    x = inputs
+    for _ in range(num_transformer_blocks):
+        x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
+
+    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
+    for dim in mlp_units:
+        x = layers.Dense(dim, activation="relu")(x)
+        x = layers.Dropout(mlp_dropout)(x)
+    outputs = layers.Dense(n_class, activation="softmax")(x)
+
+    model = Model(inputs, outputs, name="Timesweeper_Transformer_Class")
+
+    model.compile(
+        loss="categorical_crossentropy",
+        optimizer="adam",
+        metrics="accuracy",
+    )
+
+    return model
diff --git a/build/lib/timesweeper/net_hyper_tuning.py b/build/lib/timesweeper/net_hyper_tuning.py
new file mode 100644
index 0000000..16304eb
--- /dev/null
+++ b/build/lib/timesweeper/net_hyper_tuning.py
@@ -0,0 +1,308 @@
+import logging
+import os
+import pickle
+import random
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from sklearn.metrics import confusion_matrix, mean_absolute_error
+from sklearn.model_selection import train_test_split
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.layers import Conv1D, Dense, Dropout, Flatten, Input, MaxPooling1D
+from tensorflow.keras.models import Model, save_model
+from tensorflow.keras.utils import to_categorical
+from tensorflow.keras.utils import plot_model
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+from sklearn.model_selection import GridSearchCV
+from tensorflow.keras.wrappers.scikit_learn import KerasClassifier
+
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = 42
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+
+
+
+grid = ParameterGrid([{'kernel_size': [16, 32, 64, 128, 256]},
+                      {'kernel': ['rbf'], 'gamma': [1, 10]}]
+
+list(ParameterGrid(grid)) == [{'kernel': 'linear'},
+                              {'kernel': 'rbf', 'gamma': 1},
+                              {'kernel': 'rbf', 'gamma': 10}]
+
+
+
+def get_mlp_model(datadim, n_class,
+                  kernel_sizeLayerer_a=64,
+                  kernel_sizeLayerer_b=64,
+	              dropoutCovLayer_a=0.15,
+                  learnRate=0.01,
+                  add_one_layer=False,
+                  add_two_layer=False,
+                  dropoutFullLayer_a=0.2,
+                  dropoutFullLayer_b=0.2,
+                  dropoutFullLayer_c=0.1):
+    """
+        Create a CNN model based on the hypter parameter input for the whole sample
+
+    Args:
+          kernel_sizeLayerer_a (int): the kernel size of the first convolutiona layer in the first convlutional block.
+          kernel_sizeLayerer_b (int): the kernel size of the second convolutiona layer in the first convlutional block.
+          dropoutCovLayer_a (float): the dropout rate in the first convlutional block.
+          learnRate (float): the learning rate of the CNN.
+          add_one_layer (boolean): if true, add one additional convlutional block.
+          add_two_layer (boolean): if true, add a second additional convlutional block.
+          dropoutFullLayer_a (float): the dropout rate of the first layer in the fully connnected block.
+          dropoutFullLayer_b (float): the dropout rate of the second layer in the fully connnected block.
+          dropoutFullLayer_c (float): the dropout rate of the third layer in the fully connnected block.
+
+    Returns:
+        kera.classifier: a cnn fitted to the input hyperparameter.
+    """
+    ## Convolution Layer One
+    model_in = Input(datadim)
+    h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(model_in)
+    h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+    h = MaxPooling1D(pool_size=3, padding="same")(h)
+    h = Dropout(dropoutCovLayer_a)(h)
+    h = Flatten()(h)
+
+    ## Convolution Layer Two if add_one_layer
+    if add_one_layer:
+        h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(h)
+        h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+        h = MaxPooling1D(pool_size=3, padding="same")(h)
+        h = Dropout(dropoutCovLayer_a)(h)
+        h = Flatten()(h)
+
+    ## Convolution Layer Three if add_two_layer
+    if add_two_layer:
+        h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(h)
+        h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+        h = MaxPooling1D(pool_size=3, padding="same")(h)
+        h = Dropout(dropoutCovLayer_a)(h)
+        h = Flatten()(h)
+
+    ## Fully Connected Layer
+    h = Dense(264, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_a)(h)
+    h = Dense(264, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_b)(h)
+    h = Dense(128, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_c)(h)
+
+    reg_output = Dense(1, activation="relu", name="reg_output")(h)
+    class_output = Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output, reg_output], name="Timesweeper")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy", "reg_output":"mse"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy", "reg_output": "mse"},
+    )
+
+    return model
+
+
+
+def main(ua):
+    if ua.config_format == "yaml":
+        yaml_data = read_config(ua.yaml_file)
+        work_dir = yaml_data["work dir"]
+
+    elif ua.config_format == "cli":
+        work_dir = ua.work_dir
+
+    # Collect all the data
+    logger.info("Starting training process.")
+    if ua.hft:
+        type_list = ["aft", "hft"]
+    else:
+        type_list = ["aft"]
+    for data_type in type_list:
+        ids, ts_data, sweep_types, sel_coeffs = get_data(ua.training_data, data_type)
+        lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+        # Convert to numerical ohe IDs
+        num_ids = np.array([lab_dict[lab] for lab in ids])
+        ohe_ids = to_categorical(
+            np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+        )
+
+        class_weights = dict(
+            enumerate(
+                compute_class_weight(
+                    class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+                )
+            )
+        )
+        print(f"Class weights: {class_weights}")
+
+        if data_type == "aft":
+            # Needs to be in correct dims order for Conv1D layer
+            datadim = ts_data.shape[1:]
+            logger.info(
+                f"TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+            )
+        else:
+            logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+            datadim = ts_data.shape[1:]
+            logger.info(f"{len(ts_data)} samples in dataset.")
+
+        logger.info("Splitting Partition")
+        (
+            ts_train_data,
+            ts_val_data,
+            ts_test_data,
+            train_labs,
+            val_labs,
+            test_labs,
+            train_s,
+            val_s,
+            test_s,
+        ) = split_partitions(ts_data, ohe_ids, sel_coeffs)
+
+        # Time-series model training and evaluation
+        logger.info("Training time-series model.")
+
+
+        # scale data to the range of [0, 1]
+        trainData = trainData.astype("float32") / 255.0
+        testData = testData.astype("float32") / 255.0
+
+        # wrap our model into a scikit-learn compatible classifier
+        print("[INFO] initializing model...")
+
+        # define a grid of the hyperparameter search space
+        kernel_sizeLayerer_a = [16, 32, 64, 128, 256]
+        kernel_sizeLayerer_b = [16, 32, 64, 128, 256]
+        dropoutCovLayer_a= [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        # learnRate = [1e-2, 1e-3, 1e-4]
+        add_one_layer= [False, True]
+        # add_two_layer= [False, True]
+        dropoutFullLayer_a = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        dropoutFullLayer_b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        dropoutFullLayer_c = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        # batchSize = [4, 8, 16, 32]
+        # epochs = [10, 20, 30, 40]
+
+
+        # create a dictionary from the hyperparameter grid
+        grid = dict(
+        	kernel_sizeLayerer_a=kernel_sizeLayerer_a,
+        	kernel_sizeLayerer_b=kernel_sizeLayerer_b,
+        	dropoutCovLayer_a=dropoutCovLayer_a,
+        	add_one_layer=add_one_layer,
+        	dropoutFullLayer_a=dropoutFullLayer_a,
+            dropoutFullLayer_b=dropoutFullLayer_b,
+        	dropoutFullLayer_c=dropoutFullLayer_c
+        )
+
+        param_grid = ParameterGrid(grid)
+
+        for params in param_grid:
+            cur_model = get_mlp_model(**params)
+
+            # print(model.summary())
+            plot_model(
+                cur_model, to_file="model_plot.png", show_shapes=True, show_layer_names=True
+            )
+            cur_trained_model = fit_model(
+                work_dir,
+                cur_model,
+                data_type,
+                ts_train_data,
+                train_labs,
+                train_s,
+                ts_val_data,
+                val_labs,
+                val_s,
+                class_weights,
+                ua.experiment_name,
+            )
+            evaluate_model( ### need to modify to store fitted results
+                cur_trained_model,
+                ts_test_data,
+                test_labs,
+                test_s,
+                work_dir,
+                ua.experiment_name,
+                data_type,
+                lab_dict,
+            )
+
+
+        logger.info(f"SP Data shape (samples, haps): {sp_train_data.shape}")
+        # Use only the final timepoint
+        sp_train_data = np.squeeze(ts_train_data[:, -1, :])
+        sp_val_data = np.squeeze(ts_val_data[:, -1, :])
+        sp_test_data = np.squeeze(ts_test_data[:, -1, :])
+        logger.info(f"SP Data shape (samples, haps): {sp_train_data.shape}")
+
+        sp_datadim = sp_train_data.shape[-1]
+
+        for params in param_grid:
+            cur_model = model.__init__(self, **params)
+
+            sp_datadim = sp_train_data.shape[-1]
+            model = create_1Samp_model(sp_datadim, len(lab_dict))
+
+            trained_model = fit_model(
+                work_dir,
+                model,
+                data_type,
+                sp_train_data,
+                train_labs,
+                train_s,
+                sp_val_data,
+                val_labs,
+                val_s,
+                class_weights,
+                ua.experiment_name,
+            )
+            evaluate_model(
+                trained_model,
+                sp_test_data,
+                test_labs,
+                test_s,
+                work_dir,
+                ua.experiment_name,
+                data_type,
+                lab_dict,
+            )
+            )
+
+
+
+# # initialize a random search with a 3-fold cross-validation and then
+# # start the hyperparameter search process
+# print("[INFO] performing grid search...")
+# searcher = GridSearchCV(estimator=model, n_jobs=-1, cv=1
+# 	param_distributions=grid, scoring="accuracy")
+# searchResults = searcher.fit(train_data, [train_labs, train_s])
+#
+#
+# # summarize grid search information
+# bestScore = searchResults.best_score_
+# bestParams = searchResults.best_params_
+# print("[INFO] best score is {:.2f} using {}".format(bestScore,
+# 	bestParams))
+#
+# # extract the best model, make predictions on our data, and show a
+# # classification report
+# print("[INFO] evaluating the best model...")
+# bestModel = searchResults.best_estimator_
+# accuracy = bestModel.score(testData, testLabels)
+# print("accuracy: {:.2f}%".format(accuracy * 100))
diff --git a/build/lib/timesweeper/parse_slim_logs.py b/build/lib/timesweeper/parse_slim_logs.py
new file mode 100644
index 0000000..c2db182
--- /dev/null
+++ b/build/lib/timesweeper/parse_slim_logs.py
@@ -0,0 +1,109 @@
+from glob import glob
+import os
+
+import pandas as pd
+from tqdm import tqdm
+
+from timesweeper.utils.gen_utils import read_config
+
+
+def parse_cmd(cmd):
+    """Create dictionary with all information in the bash command running slim."""
+    cmd = [i for i in cmd.split() if i not in ["slim", "-d"]]
+
+    cmd_dict = {}
+    for i in cmd:
+        if "=" in i:
+            clean_i = i.split("=")
+            cmd_dict[clean_i[0].replace("'", "").replace('"', "")] = (
+                clean_i[1].replace("'", "").replace('"', "")
+            )
+
+    return cmd_dict
+
+
+def get_samp_gens(loglist):
+    """Get sampling generations."""
+    samp_entries = [i for i in loglist if "Sampling at generation" in i]
+    samp_entries = set(samp_entries)
+    samp_gens = [i.split()[-1] for i in samp_entries]
+
+    return sorted([int(i) for i in samp_gens])
+
+
+def count_restarts(loglist):
+    restarts = 0
+    for i in loglist:
+        if "RESTARTING" in i:
+            restarts += 1
+
+    return restarts
+
+
+def track_sel_freq(loglist, num_samps):
+    freq_entries = [i for i in loglist if "SEGREGATING" in i and "FIXED" not in i]
+    if len(freq_entries) > 0:
+        freqs = [float(i.split()[-1]) for i in freq_entries]
+        return freqs[len(freqs) - num_samps :]
+
+    else:
+        return [0.0] * num_samps
+
+
+def get_rep_from_filename(filename):
+    rep = filename.split("/")[-1].split(".")[0]
+    return rep
+
+
+def parse_logfile(logfile):
+    with open(logfile, "r") as ifile:
+        loglist = [i.strip() for i in ifile.readlines()]
+
+    log_dict = parse_cmd(loglist[0])
+    log_dict["sampGens"] = get_samp_gens(loglist)
+    log_dict["selAlleleFreq"] = track_sel_freq(loglist, len(log_dict["sampGens"]))
+    log_dict["numRestarts"] = count_restarts(loglist)
+    log_dict["rep"] = get_rep_from_filename(log_dict["outFileVCF"])
+    log_dict["sampOffset"] = 10000 - log_dict["sampGens"][0]
+    for i in ["outFileVCF", "outFileMS", "dumpFile"]:
+        del log_dict[i]
+
+    return log_dict
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir, schema, threads = (
+        yaml_data["work dir"],
+        ua.experiment_name,
+        ua.threads,
+    )
+
+    logfiles = glob(f"{work_dir}/logs/*/*.log", recursive=True)
+
+    log_dict_list = []
+    for l in tqdm(logfiles, desc="Parsing logs"):
+        try:
+            log_dict_list.append(parse_logfile(l))
+        except Exception as e:
+            print(l, "couldn't work")
+
+    df = pd.DataFrame(log_dict_list)
+    df = df[
+        [
+            "rep",
+            "sweep",
+            "selCoeff",
+            "sampOffset",
+            "numRestarts",
+            "numSamples",
+            "seed",
+            "physLen",
+            "sampGens",
+            "selAlleleFreq",
+        ]
+    ]
+    df.to_csv(f"{work_dir}/{schema}_params.tsv", index=False, header=True, sep="\t")
+
+    for i in logfiles:
+        os.remove(i)
diff --git a/build/lib/timesweeper/simulate_custom.py b/build/lib/timesweeper/simulate_custom.py
new file mode 100644
index 0000000..8249ba3
--- /dev/null
+++ b/build/lib/timesweeper/simulate_custom.py
@@ -0,0 +1,313 @@
+import logging
+import multiprocessing as mp
+import os
+import subprocess
+from glob import glob
+
+import numpy as np
+import yaml
+
+logging.basicConfig()
+logger = logging.getLogger("sim_custom")
+
+
+def read_config(yaml_file):
+    """Reads in the YAML config file."""
+    with open(yaml_file, "r") as infile:
+        yamldata = yaml.safe_load(infile)
+
+    return yamldata
+
+
+# Simulation
+def randomize_selCoeff_uni(lower_bound=0.000025, upper_bound=0.25):
+    """Draws selection coefficient from log uniform dist to vary selection strength."""
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+
+    return rng.uniform(lower_bound, upper_bound, 1)[0]
+
+
+def randomize_sampGens(num_timepoints, dev=50, span=200):
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+    start = round(rng.uniform(-dev, dev, 1)[0])
+    if num_timepoints == 1:
+        sampGens = [start + span]
+    else:
+        sampGens = [
+            round(i) for i in np.linspace(start, start + span + 1, num_timepoints)
+        ]
+
+    return sampGens
+
+
+def make_d_block(
+    sweep,
+    outFileVCF,
+    outFileMS,
+    dumpfile,
+    num_sample_points,
+    inds_per_tp,
+    physLen,
+    verbose=False,
+):
+    """
+    This is meant to be a very customizeable block of text for adding custom args to SLiM as constants.
+    Can add other functions to this module and call them here e.g. pulling selection coeff from a dist.
+    This block MUST INCLUDE the 'sweep' and 'outFile' params, and at the very least the outFile must be used as output for outputVCFSample.
+    Please note that when feeding strings as a constant you must escape them since this is a shell process.
+    """
+    selCoeff = randomize_selCoeff_uni()
+    sampGens = [str(i) for i in randomize_sampGens(num_sample_points)]
+
+    d_block = f"""\
+    -d "sweep='{sweep}'" \
+    -d "outFileVCF='{outFileVCF}'" \
+    -d "outFileMS='{outFileMS}'" \
+    -d "dumpFile='{dumpfile}'" \
+    -d selCoeff={selCoeff} \
+    -d sampGens='c({','.join(sampGens)})' \
+    -d numSamples={num_sample_points} \
+    -d sampleSizePerStep={inds_per_tp} \
+    -d physLen={physLen} \
+    -d seed={np.random.randint(0, 1e16)} \
+    """
+    if verbose:
+        logger.info(f"Using the following constants with SLiM: {d_block}")
+
+    return d_block
+
+
+def simulate(slim_path, d_block, slimfile, logfile):
+    cmd = " ".join([slim_path, d_block, slimfile, ">>", logfile]).replace("    ", "")
+    with open(logfile, "w") as ofile:
+        ofile.write(cmd)
+
+    try:
+
+        subprocess.run(cmd, shell=True)
+    except subprocess.CalledProcessError as e:
+        logger.error(e.output)
+
+
+# VCF Processing
+def read_multivcf(input_vcf):
+    """Reads in file and returns as list of strings."""
+    with open(input_vcf, "r") as input_file:
+        raw_lines = [i.strip() for i in input_file.readlines()]
+
+    return raw_lines
+
+
+def split_multivcf(vcf_lines, header):
+    """Splits the lines of multi-vcf file into list of vcf entries by <header> using itertools."""
+    header_idxs = [i for i in range(len(vcf_lines)) if vcf_lines[i] == header]
+
+    split_vcfs = []
+    for idx in range(len(header_idxs[:-1])):
+        split_vcfs.append(vcf_lines[header_idxs[idx] : header_idxs[idx + 1]])
+
+    split_vcfs.append(vcf_lines[header_idxs[-1] :])
+
+    return split_vcfs
+
+
+def write_vcfs(vcf_lines, vcf_dir):
+    """Writes list of vcf entries to numerically-sorted vcf files."""
+    for idx, lines in enumerate(vcf_lines):
+        with open(os.path.join(vcf_dir, f"{idx}.vcf"), "w") as outfile:
+            outfile.writelines("\n".join(lines))
+
+
+def index_vcf(vcf):
+    """
+    Indexes and sorts vcf file.
+    Commands are run separately such that processes complete before the next one starts.
+    """
+    bgzip_cmd = f"bgzip -c {vcf} > {vcf}.gz"
+    tabix_cmd = f"tabix -f -p vcf {vcf}.gz"
+    bcftools_cmd = f"bcftools sort -Ov {vcf}.gz | bgzip -f > {vcf}.sorted.gz"
+    tabix_2_cmd = f"tabix -f -p vcf {vcf}.sorted.gz"
+    subprocess.run(
+        bgzip_cmd, shell=True, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL
+    )
+    subprocess.run(
+        tabix_cmd.split(), stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL
+    )
+    subprocess.run(
+        bcftools_cmd, shell=True, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL
+    )
+    subprocess.run(
+        tabix_2_cmd.split(), stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL
+    )
+
+
+def merge_vcfs(vcf_dir):
+    num_files = len(glob(f"{vcf_dir}/*.vcf.sorted.gz"))
+    print(vcf_dir)
+    if num_files == 1:
+        cmd = f"""zcat {f"{vcf_dir}/0.vcf.sorted.gz"} > {vcf_dir}/merged.vcf"""
+
+    else:
+        cmd = f"""bcftools merge -Ov \
+                --force-samples -0 \
+                {" ".join([f"{vcf_dir}/{i}.vcf.sorted.gz" for i in range(num_files)])} > \
+                {vcf_dir}/merged.vcf \
+                """
+
+    subprocess.run(cmd, shell=True)
+
+
+def get_num_inds(vcf_file):
+    num_ind = subprocess.check_output(
+        """awk '{if ($1 == "#CHROM"){print NF-9; exit}}' """ + vcf_file,
+        shell=True,
+    )
+    return int(num_ind)
+
+
+def cleanup_intermed(vcf_dir):
+    for ifile in glob(f"{vcf_dir}/*"):
+        if "merged" not in ifile:
+            os.remove(ifile)
+
+
+def make_vcf_dir(input_vcf):
+    """Creates directory named after vcf basename."""
+    dirname = os.path.basename(input_vcf).split(".")[0]
+    dirpath = os.path.dirname(input_vcf)
+    vcf_dir = os.path.join(dirpath, dirname)
+    if os.path.exists(vcf_dir):
+        for ifile in glob(f"{vcf_dir}/*"):
+            os.remove(ifile)
+
+    os.makedirs(vcf_dir, exist_ok=True)
+
+    return vcf_dir
+
+
+def process_vcfs(input_vcf, num_tps):
+    try:
+        # Split into multiples after SLiM just concats to same file
+        raw_lines = read_multivcf(input_vcf)
+        split_lines = split_multivcf(raw_lines, "##fileformat=VCFv4.2")
+        if len(split_lines) > 0:
+            split_lines = split_lines[len(split_lines) - num_tps :]
+
+            # Creates subdir for each rep
+            vcf_dir = make_vcf_dir(input_vcf)
+            write_vcfs(split_lines, vcf_dir)
+
+            # Now index and merge
+            [index_vcf(vcf) for vcf in glob(f"{vcf_dir}/*.vcf")]
+            merge_vcfs(vcf_dir)
+            cleanup_intermed(vcf_dir)
+
+    except Exception as e:
+        print(f"[ERROR] Couldn't process {e}")
+        pass
+
+
+def simulate_prep(
+    vcf_file, num_sample_points, slimfile, slim_path, d_block, logfile, dumpFile
+):
+    simulate(slim_path, d_block, slimfile, logfile)
+    os.remove(dumpFile)
+
+    process_vcfs(vcf_file, num_sample_points)
+    os.remove(vcf_file)
+
+
+def main(ua):
+    """
+    For simulating non-stdpopsim SLiMfiles.
+    Currently only works with 1 pop models where m2 is the sweep mutation.
+    Otherwise everything else is identical to stdpopsim version, just less complex.
+
+    Generalized block of '-d' arguments to give to SLiM at the command line allow for
+    flexible script writing within the context of this wrapper. If you write your SLiM script
+    to require args set at runtime, this should be easily modifiable to do what you need and
+    get consistent results to plug into the rest of the workflow.
+
+    The two things you *will* need to specify in your '-d' args to SLiM (and somewhere in the slim script) are:
+    - sweep [str] One of "neut", "sdn", or "ssv". If you're testing only a neut/sdn model,
+        make the ssv a dummy switch for the neutral scenario.
+    - outFile [path] You will need to define this as a population outputVCFSample input, with replace=F and append=T.
+        This does *not* need to be specified by you in the custom -d block, it will be standardized to work with the rest of the pipeline using work_dir.
+        example line for slim script: `p1.outputVCFSample(sampleSizePerStep, replace=F, append=T, filePath=outFile);`
+
+    """
+    yaml_data = read_config(ua.yaml_file)
+    (
+        work_dir,
+        slim_file,
+        slim_path,
+        reps,
+        rep_range,
+        num_sample_points,
+        inds_per_tp,
+        physLen,
+    ) = (
+        yaml_data["work dir"],
+        yaml_data["slimfile"],
+        yaml_data["slim path"],
+        yaml_data["reps"],
+        ua.rep_range,
+        yaml_data["num_sample_points"],
+        yaml_data["inds_per_tp"],
+        yaml_data["physLen"],
+    )
+
+    vcf_dir = f"{work_dir}/vcfs"
+    ms_dir = f"{work_dir}/mss"
+    dumpfile_dir = f"{work_dir}/dumpfiles"
+    logfile_dir = f"{work_dir}/logs"
+
+    sweeps = ["neut", "sdn", "ssv"]
+
+    for i in [vcf_dir, dumpfile_dir, logfile_dir]:
+        for sweep in sweeps:
+            os.makedirs(f"{i}/{sweep}", exist_ok=True)
+
+    mp_args = []
+    # Inject info into SLiM script and then simulate, store params for reproducibility
+    if rep_range:  # Take priority
+        replist = range(int(rep_range[0]), int(rep_range[1]) + 1)
+    else:
+        replist = range(reps)
+
+    for rep in replist:
+        for sweep in sweeps:
+            outFileVCF = f"{vcf_dir}/{sweep}/{rep}.multivcf"
+            outFileMS = f"{ms_dir}/{sweep}/{rep}.multiMsOut"
+            dumpFile = f"{dumpfile_dir}/{sweep}/{rep}.dump"
+            logFile = f"{logfile_dir}/{sweep}/{rep}.log"
+
+            d_block = make_d_block(
+                sweep,
+                outFileVCF,
+                outFileMS,
+                dumpFile,
+                num_sample_points,
+                inds_per_tp,
+                physLen,
+                False,
+            )
+
+            mp_args.append(
+                (
+                    outFileVCF,
+                    num_sample_points,
+                    slim_file,
+                    slim_path,
+                    d_block,
+                    logFile,
+                    dumpFile,
+                )
+            )
+
+    pool = mp.Pool(processes=ua.threads)
+    pool.starmap(simulate_prep, mp_args, chunksize=1)
diff --git a/build/lib/timesweeper/simulate_stdpopsim.py b/build/lib/timesweeper/simulate_stdpopsim.py
new file mode 100644
index 0000000..715ea71
--- /dev/null
+++ b/build/lib/timesweeper/simulate_stdpopsim.py
@@ -0,0 +1,548 @@
+import argparse as ap
+import logging
+import multiprocessing as mp
+import os, sys
+import re
+import subprocess
+from itertools import cycle
+
+import numpy as np
+import pandas as pd
+
+from timesweeper.utils.gen_utils import read_config
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+logger.setLevel("INFO")
+
+
+def get_slim_code(slim_file):
+    with open(slim_file, "r") as infile:
+        raw_lines = [line.replace("\n", "") for line in infile.readlines()]
+
+    return raw_lines
+
+
+# fmt: off
+def get_slim_info(slim_lines):
+    def get_ints(lines, query):
+        """Splits a line by regex split and grabs all digits. Returns as list."""
+        # Uses max to handle max_years_b0, but might need to check if that causes other parsing issues
+        return [
+            int(s)
+            for s in re.split("\W+", [i for i in lines if query in i][0])
+            if s.isdigit()
+        ]
+
+    Q = get_ints(slim_lines, """defineConstant("Q",""")[0]
+    gen_time = get_ints(slim_lines, """defineConstant("generation_time",""")[0]
+    burn_time_mult = get_ints(slim_lines, """defineConstant("burn_in",""")[0]
+    max_years_b0 = max(get_ints(slim_lines, """defineConstant("_T",""")) #Number of years BP to sim after burn in
+    physLen = get_ints(slim_lines, """defineConstant("chromosome_length",""")[0]
+
+    pop_sizes_line = [i for i in slim_lines if "_N" in i][0]
+    pop_sizes_start = slim_lines.index(pop_sizes_line) + 2
+    # The 4,3 will need to be more flexible
+    # End of definition idx, will be end of sampling ep array
+    pop_sizes_end = (slim_lines[pop_sizes_start:].index('    defineConstant("num_epochs", length(_T));') - 2 + pop_sizes_start)  
+
+    sizes = []
+    for i in slim_lines[pop_sizes_start:pop_sizes_end]:
+        sizes.append([int(s) for s in re.split("\W+", i) if s.isdigit()])
+
+    first_biggest = max([i[0] for i in sizes])
+    burn_in_gens = first_biggest * burn_time_mult
+
+    return Q, gen_time, max_years_b0, round(burn_in_gens), physLen
+
+
+def inject_constants(raw_lines, sweep, recombRate, selCoeff, sel_gen, end_gen, mut_rate, dumpfile):
+    """Adds in sweep type, selection coefficient, and some other details."""
+    raw_lines.insert(
+        raw_lines.index("    initializeMutationRate(mutation_rate);"),
+        f"\tdefineConstant('sweep', '{sweep}');",
+    )
+    raw_lines.insert(
+        raw_lines.index("    initializeMutationRate(mutation_rate);"),
+        f"\tdefineConstant('selCoeff', Q * {selCoeff});",
+    )
+    recomb_idx = raw_lines.index("    _recombination_rates = c(")
+    raw_lines[recomb_idx] = f"    _recombination_rates = c({recombRate});"
+    raw_lines[recomb_idx+1] = "\n"
+
+    raw_lines.insert(
+        raw_lines.index("    initializeMutationRate(mutation_rate);"),
+        f"\tinitializeMutationType('m2', 0.5, 'f', selCoeff);",
+    )    
+    raw_lines.insert(
+        raw_lines.index("    initializeMutationRate(mutation_rate);"),
+        f"\tdefineConstant('dumpFile', '{dumpfile}');"
+    )
+
+    raw_lines.insert(
+        raw_lines.index("""    sim.registerLateEvent(NULL, "{dbg(self.source); end();}", G_end, G_end);"""),
+        f"""    sim.registerLateEvent(NULL, "{{dbg(self.source); checkOnSweep(); }}", {sel_gen}, {end_gen});"""
+    )
+
+    raw_lines[raw_lines.index('    defineConstant("mutation_rate", Q * 0);')] = f'    defineConstant("mutation_rate", Q * {mut_rate});'
+    
+    return raw_lines
+
+
+def sanitize_slim(raw_lines):
+    """Removes sections of code inserted by stdpopsim that aren't needed."""
+    raw_lines.pop(raw_lines.index("    sim.treeSeqOutput(trees_file);"))
+    raw_lines.pop(raw_lines.index("    initializeTreeSeq();"))
+    raw_lines.pop(
+        raw_lines.index("""            "inds=p"+pop+".sampleIndividuals("+n+"); " +""")
+    )
+    return raw_lines
+
+
+def inject_sampling(raw_lines, pop, samp_counts, gens, outfile_path):
+    """Injects the actual sampling block that outputs to VCF."""
+    samp_eps_line = [i for i in raw_lines if "sampling_episodes" in i][0]
+    # Start of sampling_episodes constant idx
+    samp_eps_start = raw_lines.index(samp_eps_line)  
+    # End of definition idx, will be end of sampling ep array
+    samp_eps_end = (raw_lines[samp_eps_start:].index("    ), c(3, 3)));") + samp_eps_start) 
+
+    new_lines = []
+    new_lines.extend(raw_lines[samp_eps_start : samp_eps_start + 1])
+    last_line_id = len(samp_counts)
+    for samps, gen, idx in zip(samp_counts, gens, range(last_line_id)):
+        if idx < last_line_id - 1:
+            new_lines.append("\t\t" + f"c({pop[1:]}, {samps}, {gen}),")
+        else:
+            new_lines.append("\t\t" + f"c({pop[1:]}, {samps}, {gen})")
+    new_lines.append("\t" + f"), c(3, {len(gens)})));")
+
+    for line in raw_lines:
+        if "treeSeqRememberIndividuals" in line:
+            raw_lines[
+                raw_lines.index(line)
+            ] = f"""\t\t\t"{pop}.outputVCFSample("+n+", replace=F, filePath='{outfile_path}', append=T);}}","""
+
+    finished_lines = raw_lines[:samp_eps_start]
+    finished_lines.extend(new_lines)
+    finished_lines.extend(raw_lines[samp_eps_end + 1 :])
+
+    return finished_lines
+
+
+def make_sel_blocks(sweep, sel_gen, pop, dumpFileName):
+    if sweep == "ssv":
+        restart_gen = sel_gen - 500
+    else:
+        restart_gen = sel_gen
+
+    intro_block = f"""
+    \n{restart_gen} late(){{
+        // save the state of the simulation 
+        print("SAVING TO " + "{dumpFileName}" + " at generation " + sim.generation);
+        sim.outputFull("{dumpFileName}");
+
+        if (sweep == "sdn")
+        {{    
+            // introduce the sweep mutation
+            target = sample({pop}.genomes, 1);
+            target.addNewDrawnMutation(m2, asInteger(chromosome_length/2));
+        }}
+    }}
+
+    """
+
+    ssv_block = f"""
+    \n{sel_gen} late(){{
+        if (sweep == "ssv")
+        {{
+            muts = sim.mutationssvype(m1);
+            if (size(muts))
+            {{
+                mut = NULL;
+                minDist = chromosome_length+1;
+                for (m in muts)
+                {{
+                    freq = sim.mutationFrequencies({pop}, m);
+                    if (freq > 0)
+                    {{
+                        dist = abs(m.position-asInteger(chromosome_length/2));
+                        if (dist < minDist)
+                        {{
+                            minDist = dist;
+                            mut = m;
+                        }}
+                    }}
+                }}
+
+                print("Chosen mut:" + mut.id);
+                mut.setMutationType(m2);
+                mut.setSelectionCoeff(selCoeff);
+                
+                print("Chose polymorphism at position " + mut.position + " and frequency " + sim.mutationFrequencies({pop}, mut) + " to become beneficial at generation " + sim.generation);
+
+            }}
+            else
+            {{
+                print("Failed to switch from neutral to beneficial at gen " + sim.generation);
+            }}
+        }}
+    }}
+    """
+
+    check_block = f"""
+    \nfunction (void)checkOnSweep(void) {{
+        m1.convertToSubstitution = F;
+        m2.convertToSubstitution = F;
+
+        if (sweep == "sdn" | (sweep == "ssv"))
+        {{
+            fixed = (sum(sim.substitutions.mutationType == m2) == 1);
+            if (fixed)
+            {{
+                print("FIXED in pop 1 at gen " + sim.generation);
+                sim.deregisterScriptBlock(self);
+            }}
+            else
+            {{
+                muts = sim.mutationssvype(m2);
+                if (size(muts) == 0)
+                {{
+                    print("LOST at gen " + sim.generation + " - RESTARTING");
+                    // Reload
+                    sim.readFromPopulationFile("{dumpFileName}");
+                
+                    // Start a newly seeded run
+                    setSeed(rdunif(1, 0, asInteger(2^32) - 1));
+                
+                    if (sweep == "sdn")
+                    {{
+                        // re-introduce the sweep mutation
+                        target = sample({pop}.genomes, 1);
+                        print("RE-INTRODUCED MUTATION at gen " + sim.generation); //" with 2Ns = " + 2*subpopSize*selCoeff);
+                        target.addNewDrawnMutation(m2, asInteger(chromosome_length/2));
+                    }}
+                }}
+            }}
+        }}
+    }}
+
+    """
+
+    all_blocks = []
+    for i in intro_block, ssv_block, check_block:
+        all_blocks.extend(i.split("\n"))
+
+    return all_blocks
+# fmt: on
+
+
+def randomize_selCoeff_loguni(lower_bound=0.005, upper_bound=0.5):
+    """Draws selection coefficient from log uniform dist to vary selection strength."""
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+    log_low = np.math.log10(lower_bound)
+    log_upper = np.math.log10(upper_bound)
+    rand_log = rng.uniform(log_low, log_upper, 1)
+
+    return 10 ** rand_log[0]
+
+
+def randomize_selCoeff_uni(lower_bound=0.0, upper_bound=0.05):
+    """Draws selection coefficient from log uniform dist to vary selection strength."""
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+
+    return rng.uniform(lower_bound, upper_bound, 1)[0]
+
+
+def randomize_recombRate(lower_bound=0, upper_bound=2e-8):
+    """Draws selection coefficient from log normal dist to vary selection strength."""
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+    recomb_rate = rng.uniform(lower_bound, upper_bound, 1)
+
+    return recomb_rate[0]
+
+
+def randomize_selTime(sel_time, stddev):
+    """Draws from uniform dist to vary the time selection is induced before sampling in gens."""
+    rng = np.random.default_rng(
+        np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+    )
+    return int(rng.uniform(sel_time - stddev, sel_time + stddev, 1)[0])
+
+
+def write_slim(finished_lines, slim_file, rep, work_dir):
+    os.makedirs(work_dir, exist_ok=True)
+    filename = os.path.basename(slim_file).split(".")[0]
+    new_file_name = os.path.join(work_dir, f"modded.{rep}.{filename}.slim")
+    with open(new_file_name, "w") as outfile:
+        for line in finished_lines:
+            outfile.write(line + "\n")
+
+    return new_file_name
+
+
+def run_slim(slimfile, slim_path):
+    cmd = f"{slim_path} {slimfile}"
+    try:
+        subprocess.check_output(cmd.split())
+    except subprocess.CalledProcessError as e:
+        logger.error(e.output)
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+
+def clean_args(ua):
+    yaml_data = read_config(ua.yaml_file)
+    (
+        slim_file,
+        reps,
+        pop,
+        sample_sizes,
+        years_sampled,
+        sel_gen,
+        sel_coeff_bounds,
+        mut_rate,
+        work_dir,
+        slim_path,
+    ) = (
+        yaml_data["slimfile"],
+        yaml_data["reps"],
+        yaml_data["pop"],
+        yaml_data["sample sizes"],
+        yaml_data["years sampled"],
+        yaml_data["selection gen"],
+        yaml_data["selection coeff bounds"],
+        yaml_data["mut rate"],
+        yaml_data["work dir"],
+        yaml_data["slim path"],
+    )
+
+    if ua.verbose:
+        logger.info(f"Number of sample sizes: {len(sample_sizes)}")
+        logger.info(f"Number of years to sample from: {len(years_sampled)}")
+
+    if len(sample_sizes) != len(years_sampled):
+        logger.error(
+            "Number of args supplied for generations and sample sizes don't match up. Please double check your values."
+        )
+        raise ValueError()
+
+    if years_sampled[0] < years_sampled[-1]:
+        logger.warning(
+            "Gens were not supplied in earliest to latest order, sorting and flipping."
+        )
+        years_sampled = sorted(years_sampled)[::-1]
+        sample_sizes = sorted(sample_sizes)[::-1]
+    else:
+        years_sampled = years_sampled
+        sample_sizes = sample_sizes
+
+    return (
+        ua.verbose,
+        ua.threads,
+        ua.rep_range,
+        work_dir,
+        slim_file,
+        reps,
+        pop,
+        sample_sizes,
+        years_sampled,
+        sel_gen,
+        sel_coeff_bounds,
+        mut_rate,
+        slim_path,
+    )
+
+
+def main(ua):
+    (
+        verbose,
+        threads,
+        rep_range,
+        work_dir,
+        slim_file,
+        reps,
+        pop,
+        sample_sizes,
+        years_sampled,
+        sel_gen,
+        sel_coeff_bounds,
+        mut_rate,
+        slim_path,
+    ) = clean_args(ua)
+
+    work_dir = work_dir
+    vcf_dir = f"{work_dir}/vcfs"
+    dumpfile_dir = f"{work_dir}/dumpfiles"
+    script_dir = f"{work_dir}/scripts"
+
+    sweeps = ["neut", "sdn", "ssv"]
+
+    for i in [vcf_dir, dumpfile_dir, script_dir]:
+        for sweep in sweeps:
+            os.makedirs(f"{i}/{sweep}", exist_ok=True)
+
+    # Inject info into SLiM script and then simulate, store params for reproducibility
+    sim_params = []
+    script_list = []
+
+    if rep_range:  # Take priority
+        replist = range(int(rep_range[0]), int(rep_range[1]))
+    else:
+        replist = range(reps)
+
+    for rep in replist:
+        for sweep in sweeps:
+
+            # Info scraping and calculations
+            raw_lines = get_slim_code(slim_file)
+            raw_lines = sanitize_slim(raw_lines)
+
+            Q, gen_time, max_years_b0, burn_in_gens, physLen = get_slim_info(raw_lines)
+
+            # Pull from variable time of selection before sampling to make more robust
+            rand_sel_gen = randomize_selTime(sel_gen, 50 / Q)
+
+            burn_in_gens = int(round(burn_in_gens / Q))
+
+            # Convert from earliest year from bp to gens
+            end_gen = int(round((max_years_b0) / gen_time / Q))
+
+            # Written out each step for clarity, sdn to keep track of otherwise
+            # Find the earliest time before present, convert to useable times
+            furthest_from_pres = max(years_sampled)
+            abs_year_beg = max_years_b0 - furthest_from_pres
+
+            sel_gen_time = (
+                int(((abs_year_beg / gen_time) - rand_sel_gen) / Q)
+            ) + burn_in_gens
+
+            if sel_coeff_bounds[0] == sel_coeff_bounds[1]:
+                sel_coeff = sel_coeff_bounds[0] * Q
+            else:
+                sel_coeff = randomize_selCoeff_uni(*sel_coeff_bounds) * Q
+
+            recombRate = randomize_recombRate()
+
+            # Logger vars
+            if verbose:
+                logger.info("Timesweeper SLiM Injection")
+                logger.info(f"Q Scaling Value: {Q}")
+                logger.info(f"Gen Time: {gen_time}")
+                logger.info(f"Simulated Chrom Length: {physLen}")
+                logger.info(f"Burn in years: {burn_in_gens * gen_time}")
+                logger.info(f"Burn in gens: {burn_in_gens}")
+                logger.info(f"Number Years Simulated (post-burn): {max_years_b0}")
+                logger.info(f"Number Gens Simulated (post-burn): {end_gen}")
+                logger.info(
+                    f"Number Years Simulated (inc. burn): {max_years_b0 + (burn_in_gens * gen_time)}"
+                )
+                logger.info(
+                    f"Number gens simulated (inc. burn): {end_gen + burn_in_gens}"
+                )
+                logger.info(f"Selection type: {sweep}")
+                logger.info(f"Selection coeff: {sel_coeff}")
+                logger.info(f"Recomb Rate: {recombRate}")
+                logger.info(f"Selection start gen: {sel_gen_time}")
+                logger.info(f"Number of timepoints: {len(years_sampled)}")
+                logger.info(
+                    f"""Sample sizes (individuals): 
+                    {" ".join([str(i) for i in sample_sizes])}"""
+                )
+                logger.info(
+                    f"Years before present (1950) sampled: {' '.join([str(i) for i in years_sampled])}"
+                )
+                logger.info(
+                    f"Gens before present (1950) sampled: {' '.join([str(int(i / gen_time / Q)) for i in years_sampled])}",
+                )
+
+            sim_params.append(
+                (
+                    sweep,
+                    rep,
+                    Q,
+                    gen_time,
+                    physLen,
+                    burn_in_gens,
+                    max_years_b0 + (burn_in_gens * gen_time),
+                    sel_coeff,
+                    sel_gen_time,
+                    years_sampled,
+                    sample_sizes,
+                )
+            )
+
+            dumpfile = f"{dumpfile_dir}/{sweep}/{rep}.dump"
+
+            # Injection
+            prepped_lines = inject_constants(
+                raw_lines,
+                sweep,
+                recombRate,
+                sel_coeff,
+                sel_gen_time,
+                end_gen + burn_in_gens,
+                mut_rate,
+                dumpfile,
+            )
+
+            sampling_lines = inject_sampling(
+                prepped_lines,
+                pop,
+                sample_sizes,
+                years_sampled,
+                f"{vcf_dir}/{sweep}/{rep}.multivcf",
+            )
+
+            selection_lines = make_sel_blocks(sweep, sel_gen_time, pop, dumpfile)
+            finished_lines = []
+            finished_lines.extend(sampling_lines)
+            finished_lines.extend(selection_lines)
+
+            script_path = write_slim(
+                finished_lines, slim_file, rep, f"{script_dir}/{sweep}"
+            )
+            script_list.append(script_path)
+
+    print(f"Reps simulated: {replist}")
+
+    for script in script_list:
+        run_slim(script, slim_path)
+
+    # Cleanup
+    for rep in replist:
+        for sweep in sweeps:
+            dumpFile = f"{dumpfile_dir}/{sweep}/{rep}.dump"
+            os.remove(dumpFile)
+
+    # Save params
+    if not os.path.exists(f"{work_dir}/params"):
+        os.makedirs(f"{work_dir}/params")
+    params_df = pd.DataFrame(
+        sim_params,
+        columns=[
+            "sweep",
+            "rep",
+            "Q",
+            "gen_time",
+            "chrom_len",
+            "burn_in_gens",
+            "num_gens_simmed",
+            "sel_coeff",
+            "sel_gen",
+            "years_bp_sampled",
+            "samp_sizes",
+        ],
+    )
+    params_df.to_csv(
+        f"{work_dir}/params/sim_{'_'.join(rep_range)}_params.tsv", sep="\t", index=False
+    )
+
+    logger.info(f"Simulations finished, parameters saved to {work_dir}/sim_params.csv.")
diff --git a/build/lib/timesweeper/tp1_model.py b/build/lib/timesweeper/tp1_model.py
new file mode 100644
index 0000000..d82c0bd
--- /dev/null
+++ b/build/lib/timesweeper/tp1_model.py
@@ -0,0 +1,743 @@
+"""
+This is a modified version of the `train_nets.py` module in the core Timesweeper package. 
+Modifications include:
+- Fully connected network instead of 1DCNN
+"""
+
+import logging
+import os
+import pickle
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    id_list = []
+    rep_list = []
+    data_list = []
+    sel_coeffs = []
+    sweep_types = []
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+    for sweep in pikl_dict.keys():
+        sweep_types.append(sweep)
+        for rep in pikl_dict[sweep].keys():
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+            id_list.append(sweep)
+            rep_list.append(rep)
+            sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+
+    return (
+        id_list,
+        rep_list,
+        np.stack(data_list),
+        sweep_types,
+        np.array(sel_coeffs).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_s,
+    val_data,
+    val_s,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    ts_data = np.squeeze(ts_data)
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_1tp_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_1tp_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = False
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/build/lib/timesweeper/train_nets.py b/build/lib/timesweeper/train_nets.py
new file mode 100644
index 0000000..4d3a46e
--- /dev/null
+++ b/build/lib/timesweeper/train_nets.py
@@ -0,0 +1,737 @@
+import logging
+import os
+import pickle
+import random
+import sys
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    id_list = []
+    rep_list = []
+    data_list = []
+    sel_coeffs = []
+    sweep_types = []
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+    for sweep in pikl_dict.keys():
+        sweep_types.append(sweep)
+        for rep in pikl_dict[sweep].keys():
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+            
+            id_list.append(sweep)
+            rep_list.append(rep)
+            sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+
+    return (
+        id_list,
+        rep_list,
+        np.stack(data_list),
+        sweep_types,
+        np.array(sel_coeffs).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir, model, data_type, train_data, train_s, val_data, val_s, experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "2dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "2dcnn":
+        class_model = models.create_2D_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_2D_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "chonk":
+        class_model = models.create_big_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_big_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "rnn":
+        class_model = models.create_rnn_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_rnn_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = True
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/build/lib/timesweeper/train_nets_shic.py b/build/lib/timesweeper/train_nets_shic.py
new file mode 100644
index 0000000..f63a590
--- /dev/null
+++ b/build/lib/timesweeper/train_nets_shic.py
@@ -0,0 +1,720 @@
+import logging
+import os
+import pickle
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+
+    return (
+        pikl_dict["sweep"],
+        pikl_dict["rep"],
+        np.stack(pikl_dict["data"]),
+        ["neut", "sdn", "ssv"],
+        np.array(pikl_dict["selcoeff"]).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_s,
+    val_data,
+    val_s,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = True
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/build/lib/timesweeper/train_shoulder_nets.py b/build/lib/timesweeper/train_shoulder_nets.py
new file mode 100644
index 0000000..86d15ee
--- /dev/null
+++ b/build/lib/timesweeper/train_shoulder_nets.py
@@ -0,0 +1,715 @@
+import logging
+import os
+import pickle
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type, win_size):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    id_list = []
+    rep_list = []
+    data_list = []
+    sel_coeffs = []
+    sweep_types = []
+    offsets = []
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+    for sweep in pikl_dict.keys():
+        sweep_types.append(sweep)
+        for rep in pikl_dict[sweep].keys():
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+
+            if sweep != "neut":
+                if abs(pikl_dict[sweep][rep]["center_offset"]) > 500 + win_size:
+                    id_list.append("neut")
+                else:
+                    id_list.append(sweep)
+
+            else:
+                id_list.append(sweep)
+
+            rep_list.append(rep)
+            sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+            offsets.append(pikl_dict[sweep][rep]["center_offset"])
+
+    return (
+        id_list,
+        rep_list,
+        np.stack(data_list),
+        sweep_types,
+        np.array(sel_coeffs).reshape(-1, 1),
+        np.array(offsets).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps, offsets):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+        train_offsets,
+        val_offsets,
+    ) = train_test_split(
+        data, labs, sel_coeffs, reps, offsets, stratify=labs, test_size=0.3
+    )
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+        val_offsets,
+        test_offsets,
+    ) = train_test_split(
+        val_data,
+        val_labs,
+        val_s,
+        val_reps,
+        val_offsets,
+        stratify=val_labs,
+        test_size=0.5,
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+        train_offsets,
+        val_offsets,
+        test_offsets,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_s,
+    val_data,
+    val_s,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    test_offsets,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "offset": test_offsets.flatten(),
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    test_offsets,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "offset": test_offsets.flatten(),
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info(
+        "Starting training process, shoulders are included in the training set."
+    )
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs, raw_offsets = get_data(
+        ua.training_data, ua.data_type, yaml_data["win_size"]
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        (
+            ts_data,
+            _,
+            ohe_ids,
+            _,
+            sel_coeffs,
+            _,
+            reps,
+            _,
+            offsets,
+            _,
+        ) = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            raw_offsets,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+        offsets = raw_offsets
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+        train_offsets,
+        val_offsets,
+        test_offsets,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps, offsets)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            test_offsets,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s)
+                vvals = -np.log10(val_s)
+                tevals = -np.log10(test_s)
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s)
+                trvals = mm_scaler.transform(train_s)
+                vvals = mm_scaler.transform(val_s)
+                tevals = mm_scaler.transform(test_s)
+            else:
+                mm_scaler = None
+                trvals = train_s
+                vvals = val_s
+                tevals = test_s
+
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape} {ts_val_data.shape} {ts_test_data.shape}"
+            )
+
+            plot = True
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s,
+                    ts_train_data[:, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[:, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data,
+                trvals,
+                ts_val_data,
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data,
+                test_labs,
+                tevals,
+                test_reps,
+                test_offsets,
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/meta.yaml b/meta.yaml
deleted file mode 100644
index 35e2d14..0000000
--- a/meta.yaml
+++ /dev/null
@@ -1,30 +0,0 @@
-package:
-  name: timesweeper
-  version: "1.0.0"
-
-source:
-  git_url: https://github.com/SchriderLab/Timesweeper.git
-
-requirements:
-  build:
-    - python
-    - setuptools
-
-  run:
-   - numpy
-   - scipy
-   - matplotlib
-   - pandas
-   - scikit-allel
-   - scikit-learn
-   - tensorflow
-   - keras
-   - tqdm
-   - slim
-   - samtools
-   - bcftools
-   
-
-about:
-  home: https://github.com/SchriderLab/Timesweeper
-  license: MIT
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 9048b9a..e130edd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,8 +1,64 @@
-numpy
-pandas
-tensorflow
-scikit-learn
-scikit-allel
-tqdm
-matplotlib
-scipy
\ No newline at end of file
+absl-py==1.4.0
+astunparse==1.6.3
+cachetools==5.3.0
+certifi==2022.12.7
+charset-normalizer==3.1.0
+click==8.1.3
+cloudpickle==2.2.1
+contourpy==1.0.7
+cycler==0.11.0
+dask==2023.3.0
+flatbuffers==23.3.3
+fonttools==4.39.0
+fsspec==2023.3.0
+gast==0.4.0
+google-auth==2.16.2
+google-auth-oauthlib==0.4.6
+google-pasta==0.2.0
+grpcio==1.51.3
+h5py==3.8.0
+idna==3.4
+jax==0.4.5
+joblib==1.2.0
+keras==2.12.0rc1
+kiwisolver==1.4.4
+libclang==15.0.6.1
+locket==1.0.0
+Markdown==3.4.1
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+numpy==1.23.5
+oauthlib==3.2.2
+opt-einsum==3.3.0
+packaging==23.0
+pandas==1.5.3
+partd==1.3.0
+Pillow==9.4.0
+protobuf==4.22.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pyparsing==3.0.9
+python-dateutil==2.8.2
+pytz==2022.7.1
+PyYAML==6.0
+requests==2.28.2
+requests-oauthlib==1.3.1
+rsa==4.9
+scikit-allel==1.3.5
+scikit-learn==1.2.1
+scipy==1.10.1
+six==1.16.0
+tensorboard==2.12.0
+tensorboard-data-server==0.7.0
+tensorboard-plugin-wit==1.8.1
+tensorflow==2.12.0rc0
+tensorflow-estimator==2.12.0rc0
+tensorflow-io-gcs-filesystem==0.31.0
+termcolor==2.2.0
+threadpoolctl==3.1.0
+toolz==0.12.0
+tqdm==4.65.0
+typing_extensions==4.5.0
+urllib3==1.26.14
+Werkzeug==2.2.3
+wrapt==1.15.0
diff --git a/setup.cfg b/setup.cfg
index e3f4e5d..2a818d6 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,14 +1,13 @@
 [metadata]
 name = Timesweeper
-version = 1.1.0
+version = 1.1.2
 author = Logan Whitehouse
 author_email = lswhiteh@unc.edu
 description = A tool for detecting positive selective sweeps using time-series genomcis data.
-long_description = file: README.md, LICENSE
-url = https://github.com/SchriderLab/timeSeriesSweeps
+url = https://github.com/SchriderLab/Timesweeper
 project_urls = 
-    Issues = https://github.com/SchriderLab/timeSeriesSweeps/issues
-    GitHub = https://github.com/SchriderLab/timeSeriesSweeps
+    Issues = https://github.com/SchriderLab/Timesweeper/issues
+    GitHub = https://github.com/SchriderLab/Timesweeper
 classifiers = 
     Programming Language :: Python :: 3
     Operating System :: OS Independent
@@ -22,7 +21,15 @@ license = MIT
 python_requires = >=3.6
 include_package_data = True
 packages = timesweeper
-
+install_requires = 
+    numpy
+    pandas
+    tensorflow
+    scikit-learn
+    scikit-allel
+    tqdm
+    matplotlib
+    scipy
 
 [options.entry_points]
 console_scripts =
diff --git a/setup.py b/setup.py
index 7f1a176..4506e8c 100644
--- a/setup.py
+++ b/setup.py
@@ -1,4 +1,7 @@
 from setuptools import setup
 
 if __name__ == "__main__":
-    setup()
+    with open("README.md", "r") as readme:
+        readme_txt = readme.read()
+    
+    setup(long_description=readme_txt, long_description_content_type='text/markdown')
diff --git a/timesweeper/cli.py b/timesweeper/cli.py
index 779bf87..4699bc6 100755
--- a/timesweeper/cli.py
+++ b/timesweeper/cli.py
@@ -11,7 +11,7 @@ def ts_main():
     # simulate_stdpopsim.py
     sim_s_parser = subparsers.add_parser(
         name="sim_stdpopsim",
-        description="Injects time-series sampling into stdpopsim SLiM output script.",
+        help="Injects time-series sampling into stdpopsim SLiM output script.",
     )
     sim_s_parser.add_argument(
         "-v",
@@ -46,7 +46,7 @@ def ts_main():
     # simulate_custom.py
     sim_c_parser = subparsers.add_parser(
         name="sim_custom",
-        description="Simulates selection for training Timesweeper using a pre-made SLiM script.",
+        help="Simulates selection for training Timesweeper using a pre-made SLiM script.",
     )
     sim_c_parser.add_argument(
         "--threads",
@@ -74,43 +74,29 @@ def ts_main():
     )
 
     # process_vcfs.py
-    vcfproc_parser = subparsers.add_parser(
+    process_vcf_parser = subparsers.add_parser(
         name="process",
-        description="Splits and re-merges VCF files to prepare for fast feature creation.",
+        help="Module for splitting multivcfs (vertically concatenated vcfs) into merged (horizontally concatenated vcfs) if simulating without the sim module.",
     )
-    vcfproc_parser.add_argument(
-        "--vcf-header",
-        required=False,
-        type=str,
-        default="##fileformat=VCFv4.2",
-        dest="vcf_header",
-        help="String that tops VCF header, used to split entries to new files.",
+    process_vcf_parser.add_argument(
+        "-i",
+        "--in-dir",
+        dest="in_dir",
+        help="Top-level directory containing subidrectories labeled with the names of each scenario and then replicate numbers inside those containing each multivcf.",
+        required=True,
     )
-    vcfproc_parser.add_argument(
+    process_vcf_parser.add_argument(
+        "-t",
         "--threads",
-        required=False,
-        type=int,
-        default=mp.cpu_count() - 1,
         dest="threads",
-        help="Number of processes to parallelize across.",
-    )
-    vcfproc_parser.add_argument(
-        "-y",
-        "--yaml",
-        metavar="YAML_CONFIG",
+        help="Threads to use for multiprocessing.",
         required=True,
-        dest="yaml_file",
-        help="YAML config file with all required options defined.",
-    )
-
+    )   
+     
     # make_training_features.py
     mtf_parser = subparsers.add_parser(
         name="condense",
-        description="Creates training data from simulated merged vcfs after process_vcfs.py has been run.",
-    )
-    mtf_parser = subparsers.add_parser(
-        name="condense",
-        description="Creates training data from simulated merged vcfs after process_vcfs.py has been run.",
+        help="Creates training data from simulated merged vcfs after process_vcfs.py has been run.",
     )
     mtf_parser.add_argument(
         "--threads",
@@ -129,6 +115,27 @@ def ts_main():
         dest="outfile",
         help="Pickle file to dump dictionaries with training data to. Should probably end with .pkl.",
     )
+    mtf_parser.add_argument(
+        "--subsample-inds",
+        required=False,
+        type=int,
+        dest="subsample_inds",
+        help="Number of individuals to subsample if using a larger simulation than needed for efficiency. NOTE: If you use this, keep the sample_sizes entry in the yaml config identical to the original simulations. This is mostly for the paper experiments and as such isn't very cleanly implemented.",
+    )
+    mtf_parser.add_argument(
+        "--subsample-tps",
+        required=False,
+        type=int,
+        dest="subsample_tps",
+        help="Number of timepoints to subsample if using a larger simulation than needed for efficiency. NOTE: If you use this, keep the sample_sizes entry in the yaml config identical to the original simulations. This is mostly for the paper experiments and as such isn't very cleanly implemented.",
+    )
+    mtf_parser.add_argument(
+        "--og-tps",
+        required=False,
+        type=int,
+        dest="og_tps",
+        help="Number of timepoints taken in the original data to be subsetted if using --subample-tps.",
+    )
     mtf_parser.add_argument(
         "-m",
         "--missingness",
@@ -139,7 +146,6 @@ def ts_main():
         default=0.0,
         help="Missingness rate in range of [0,1], used as the parameter of a binomial distribution for randomly removing known values.",
     )
-    """
     mtf_parser.add_argument(
         "-f",
         "--freq-increase-threshold",
@@ -150,7 +156,13 @@ def ts_main():
         default=0.0,
         help="If given, only include sim replicates where the sweep site has a minimum increase of <freq_inc_thr> from the first timepoint to the last.",
     )
-    """
+    mtf_parser.add_argument(
+        "-a",
+        "--allow-shoulders",
+        dest="allow_shoulders",
+        action="store_true",
+        help="1/3 of samples from sweep classes will be offset to be used as neutral shoulders.",
+    )
     mtf_parser.add_argument(
         "--hft",
         required=False,
@@ -178,10 +190,11 @@ def ts_main():
         dest="yaml_file",
         help="YAML config file with all required options defined.",
     )
-    # nets.py
+
+    # train_nets.py
     nets_parser = subparsers.add_parser(
         name="train",
-        description="Handler script for neural network training and prediction for TimeSweeper Package.\
+        help="Handler script for neural network training and prediction for TimeSweeper Package.\
             Will train two models: one for the series of timepoints generated using the hfs vectors over a timepoint and one ",
     )
     nets_parser.add_argument(
@@ -194,11 +207,21 @@ def ts_main():
         help="Pickle file containing data formatted with make_training_features.py.",
     )
     nets_parser.add_argument(
-        "--hft",
+        "-d",
+        "--data-type",
+        required=True,
+        metavar="DATA_TYPE",
+        dest="data_type",
+        help="AFT or HFT data preparation.",
+    )
+    nets_parser.add_argument(
+        "-s",
+        "--subsample-amount",
+        metavar="SUBSAMPLE_AMOUNT",
+        dest="subsample_amount",
+        type=int,
         required=False,
-        action="store_true",
-        dest="hft",
-        help="Whether to calculate HFT alongside AFT. Computationally more expensive.",
+        help="Amount of data to subsample for each class to test for sample size effects.",
     )
     nets_parser.add_argument(
         "-n",
@@ -218,11 +241,18 @@ def ts_main():
         dest="yaml_file",
         help="YAML config file with all required options defined.",
     )
+    nets_parser.add_argument(
+        "--single-tp",
+        dest="single_tp",
+        required=False,
+        action="store_true",
+        help="Whether to use the tp1_model module for a special case experiment.",
+    )
 
     # find_sweeps.py
     sweeps_parser = subparsers.add_parser(
         name="detect",
-        description="Module for iterating across windows in a time-series vcf file and predicting whether a sweep is present at each snp-centralized window.",
+        help="Module for iterating across windows in a time-series vcf file and predicting whether a sweep is present at each snp-centralized window.",
     )
     sweeps_parser.add_argument(
         "-i",
@@ -241,22 +271,34 @@ def ts_main():
         required=False,
     )
     sweeps_parser.add_argument(
-        "--aft-model",
-        dest="aft_model",
-        help="Path to Keras2-style saved model to load for aft prediction.",
+        "--aft-class-model",
+        dest="aft_class_model",
+        help="Path to Keras2-style saved model to load for classification aft prediction.",
+        required=True,
+    )
+    sweeps_parser.add_argument(
+        "--aft-reg-model",
+        dest="aft_reg_model",
+        help="Path to Keras2-style saved model to load for regression aft prediction. Either SDN or SSV work, both will be loaded.",
         required=True,
     )
     sweeps_parser.add_argument(
-        "--hft-model",
-        dest="hft_model",
-        help="Path to Keras2-style saved model to load for hft prediction.",
+        "--hft-class-model",
+        dest="hft_class_model",
+        help="Path to Keras2-style saved model to load for classification hft prediction.",
         required=False,
     )
+    sweeps_parser.add_argument(
+        "--hft-reg-model",
+        dest="hft_reg_model",
+        help="Path to Keras2-style saved model to load for regression hft prediction.",
+        required=True,
+    )
     sweeps_parser.add_argument(
         "-o",
-        "--out-dir",
-        dest="outdir",
-        help="Directory to write output to.",
+        "--outfile-prefix",
+        dest="outfile",
+        help="Prefix to use for results csvs. Will have '_[aft/hft].csv' appended to it.",
         required=True,
     )
     sweeps_parser.add_argument(
@@ -267,11 +309,19 @@ def ts_main():
         dest="yaml_file",
         help="YAML config file with all required options defined.",
     )
-    
-    #plot_training_data.py
+    sweeps_parser.add_argument(
+        "-s",
+        "--scalar",
+        metavar="SCALAR",
+        required=True,
+        dest="scalar",
+        help="Minmax scalar saved as a pickle file during training.",
+    )
+
+    # plot_training_data.py
     input_plot_parser = subparsers.add_parser(
         name="plot_training",
-        description="Plots central SNPs from simulations to visually inspect mean trends over replicates."
+        help="Plots central SNPs from simulations to visually inspect mean trends over replicates.",
     )
     input_plot_parser.add_argument(
         "-i",
@@ -312,7 +362,7 @@ def ts_main():
 
     freq_plot_parser = subparsers.add_parser(
         name="plot_freqs",
-        description="Create a bedfile of major and minor allele frequency changes over time."
+        help="Create a bedfile of major and minor allele frequency changes over time.",
     )
     freq_plot_parser.add_argument(
         "-i",
@@ -321,7 +371,7 @@ def ts_main():
         metavar="INPUT VCF FILE",
         type=str,
         required=True,
-        help="Merged time-series VCF file to pull SNPs and frequencies from."
+        help="Merged time-series VCF file to pull SNPs and frequencies from.",
     )
     freq_plot_parser.add_argument(
         "-o",
@@ -345,41 +395,121 @@ def ts_main():
         dest="yaml_file",
         help="YAML config file with all required options defined.",
     )
+
+    # parse_sim_logs.py
+    summarize_parser = subparsers.add_parser(
+        name="summarize",
+        help="Creates a CSV of data parsed from slim log files.",
+    )
+    summarize_parser.add_argument(
+        "--threads",
+        required=False,
+        type=int,
+        default=16,
+        dest="threads",
+        help="Number of processes to parallelize across. Defaults to all.",
+    )
+    summarize_parser.add_argument(
+        "-n",
+        "--experiment-name",
+        metavar="EXPERIMENT_NAME",
+        dest="experiment_name",
+        type=str,
+        required=False,
+        default="ts_experiment",
+        help="Identifier for the experiment used to generate the data. Optional, but helpful in differentiating runs.",
+    )
+    summarize_parser.add_argument(
+        "-y",
+        "--yaml",
+        metavar="YAML_CONFIG",
+        required=True,
+        dest="yaml_file",
+        help="YAML config file with all required options defined.",
+    )
+
+    # merge_locs_acc.py
+    merge_parser = subparsers.add_parser(
+        name="merge_logs",
+        help="Merges the summary TSV from SLiM logs with test data predictions.",
+    )
+    merge_parser.add_argument(
+        "-s",
+        "--summary-tsv",
+        dest="summary_tsv",
+        help="TSV created using `parse_slim_logs.py`.",
+        required=True,
+    )
+    merge_parser.add_argument(
+        "-i",
+        "--input-test-file",
+        dest="test_file",
+        help="Files to merge ",
+        required=True,
+    )
+    merge_parser.add_argument(
+        "-o",
+        "--output-dir",
+        metavar="OUTPUT DIR",
+        dest="output_dir",
+        required=True,
+        type=str,
+        help="Directory to write merged results and plots to.",
+    )
+
     ua = agp.parse_args()
 
-    #fmt: off
+    # fmt: off
     if ua.mode == "sim_stdpopsim":
-        from . import simulate_stdpopsim
+        from timesweeper import simulate_stdpopsim
         simulate_stdpopsim.main(ua)
 
     elif ua.mode == "sim_custom":
-        from . import simulate_custom
+        from timesweeper import simulate_custom
         simulate_custom.main(ua)
-
+    
     elif ua.mode == "process":
-        from . import process_vcfs
-        process_vcfs.main(ua)
-        
+        from timesweeper import process_vcfs
+        process_vcfs.main(ua)    
+
     elif ua.mode == "condense":
-        from . import make_training_features
+        from timesweeper import make_training_features
         make_training_features.main(ua)    
 
     elif ua.mode == "train":
-        from . import nets
-        nets.main(ua)  
+        if ua.single_tp:
+            from timesweeper import tp1_model
+            tp1_model.main(ua)   
+        elif "shoulder" in str(ua.experiment_name).lower():
+            from timesweeper import train_shoulder_nets
+            train_shoulder_nets.main(ua)  
+        elif "shic" in str(ua.experiment_name).lower():
+            from timesweeper import train_nets_shic
+            train_nets_shic.main(ua)
+        else:
+            from timesweeper import train_nets
+            train_nets.main(ua)
         
     elif ua.mode == "detect":
-        from . import find_sweeps_vcf as find_sweeps_vcf
+        from timesweeper import find_sweeps_vcf as find_sweeps_vcf
         find_sweeps_vcf.main(ua)   
 
     elif ua.mode == "plot_training":
-        from .plotting import plot_training_data as plot_training
+        from timesweeper.plotting import plot_training_data as plot_training
         plot_training.main(ua)   
 
     elif ua.mode == "plot_freqs":
-        from .plotting import create_freq_track as cf
+        from timesweeper.plotting import create_freq_track as cf
         cf.main(ua)
-
+        
+    elif ua.mode == "summarize":
+        from timesweeper import parse_slim_logs as psl
+        psl.main(ua)
+    
+    elif ua.mode == "merge_logs":
+        from timesweeper import merge_logs_acc as mla
+        mla.main(ua)
+        
     elif ua.mode == None:
         agp.print_help()
 
diff --git a/timesweeper/constant_population.slim b/timesweeper/constant_population.slim
new file mode 100644
index 0000000..9c00697
--- /dev/null
+++ b/timesweeper/constant_population.slim
@@ -0,0 +1,161 @@
+initialize()
+{
+    //remove these constants from your version since this stuff should all be set on the command line
+    defineConstant("subpopSize", 500);
+    defineConstant("mutLoc", asInteger(physLen/2));
+    defineConstant("sweepDuration", 200); //NOTE: everything will work so long as samplingInterval*numSamples is less than this value; so whatever the max value of that product will be of all the combinations you are using, just make sure this is >= to that
+    initializeMutationRate(1e-7);
+    initializeMutationType("m1", 0.5, "f", 0.0);
+    initializeGenomicElementType("g1", m1, 1.0);
+    initializeGenomicElement(g1, 0, physLen-1);
+    initializeRecombinationRate(1e-7);
+    
+    defineConstant("burnTime", 20*subpopSize);
+    defineConstant("mutTime", burnTime);
+    defineConstant("ssvTime", mutTime); 
+    defineConstant("simEndTime", burnTime+sampGens[size(sampGens)-1]);
+    defineConstant("samplingTime", burnTime+sampGens[0]); 
+    // introduced mutation's selection coefficient
+    initializeMutationType("m2", 0.5, "f", selCoeff);
+
+    defineConstant("restartTime", burnTime-subpopSize);
+    
+}
+
+s1 100 late()
+{
+    // save the state of the simulation 
+    cat("SAVING TO " + dumpFile + " at generation " + sim.generation + "\n");
+    sim.outputFull(dumpFile);  
+}
+
+s2 1000 late()
+{  
+    if (sweep == "sdn")
+    {
+        // introduce the sweep mutation
+        target = sample(p1.genomes, 1);
+        cat("INTRODUCED MUTATION at gen " + sim.generation + " with 2Ns = " + 2*subpopSize*selCoeff + "\n");
+        target.addNewDrawnMutation(m2, mutLoc);
+    }
+    
+    if (sweep == "ssv")
+    {
+        muts = sim.mutationsOfType(m1);
+        if (size(muts))
+        {
+            mut = NULL;
+            minDist = physLen+1;
+            for (m in muts)
+            {
+                dist = abs(m.position-mutLoc);
+                if (dist < minDist)
+                {
+                    minDist = dist;
+                    mut = m;
+                }
+            }
+            
+            cat("chosen mut:" + mut.id + "\n");
+            mut.setMutationType(m2);
+            mut.setSelectionCoeff(selCoeff);
+            cat(simID + ": chose polymorphism at position " + mut.position + " and frequency " + sim.mutationFrequencies(p1, mut) + " to become beneficial at generation " + sim.generation + "\n\n");
+
+        }
+        else
+        {
+            cat(simID + ": failed to switch from neutral to beneficial at gen " + sim.generation + "\n");
+        }
+    }
+}
+
+s3 1000:10000 late()
+{
+    if (sweep == "sdn" | (sweep == "ssv" & sim.generation > ssvTime))
+    {
+        fixed = (sum(sim.substitutions.mutationType == m2) == 1);
+        if (fixed)
+        {
+            sim.deregisterScriptBlock(self);
+        }
+        else
+        {
+            muts = sim.mutationsOfType(m2);
+            if (size(muts) == 0)
+            {
+                cat(simID + ": LOST at gen " + sim.generation + "\n");
+
+                // Go back to generation 1000 or ssvGen
+                sim.readFromPopulationFile(dumpFile);
+
+                // Start a newly seeded run
+                setSeed(rdunif(1, 0, asInteger(2^32) - 1));
+                
+                cat("RESTARTING WITH SEED: " + getSeed() + "\n");
+            }
+        }
+    }
+}
+
+s4 10000:15000 late()
+{
+    m1.convertToSubstitution = F;
+    m2.convertToSubstitution = F;
+    
+    if (any(mutTime+sampGens==sim.generation))
+    {
+        if (sweep == "sdn" | sweep == "ssv")
+        {   
+            if (sim.generation > mutTime)
+            {
+                muts = sim.mutationsOfType(m2);
+                if (size(muts) > 0)
+                {
+                    freq1 = sim.mutationFrequencies(p1, muts)[0];
+                    cat("SEGREGATING at " + freq1 + "\n");
+                }
+                else
+                {
+                    fixed = (sum(sim.substitutions.mutationType == m2) == 1);
+                    cat("NO LONGER SEGREGATING at generation " + sim.generation + "; mut was " + ifelse(fixed, "FIXED\n", "LOST\n") + "\n");
+                }
+            }
+        }
+
+        cat("Sampling at generation " + sim.generation + "\n");
+        p1.outputVCFSample(sampleSizePerStep, replace=F, filePath=outFileVCF, append=T);
+        //p1.outputMSSample(sampleSizePerStep*2, replace=F, filePath=outFileMS, append=T);
+
+        cat("Done emitting sample\n");
+
+        if (sim.generation == simEndTime)
+        {
+            p1.outputVCFSample(size(p1.individuals), replace=F, filePath=outFileVCF+".final", append=F);
+        }
+    }
+}
+
+s5 15000 late()
+{
+    sim.simulationFinished();
+}
+
+
+1 early() {
+    // save this run's identifier, used to save and restore
+    defineConstant("simID", getSeed());
+    
+    // set up our subpop
+    sim.addSubpop("p1", subpopSize);
+    
+    //schedule our events
+
+
+    sim.rescheduleScriptBlock(s1, start=restartTime, end=restartTime); //  set far restart time
+    sim.rescheduleScriptBlock(s2, start=mutTime, end=mutTime); // introducing mut
+    sim.rescheduleScriptBlock(s3, start=mutTime, end=simEndTime); // checking on muts
+    sim.rescheduleScriptBlock(s4, start=samplingTime, end=simEndTime); // sampling
+    sim.rescheduleScriptBlock(s5, start=simEndTime, end=simEndTime); // wrapping up
+
+    
+}
\ No newline at end of file
diff --git a/timesweeper/convert_syncfiles.py b/timesweeper/convert_syncfiles.py
new file mode 100644
index 0000000..731fe6a
--- /dev/null
+++ b/timesweeper/convert_syncfiles.py
@@ -0,0 +1,176 @@
+import sys
+import numpy as np
+import argparse
+
+
+
+def getRepAndGen(headerEntry):
+    genInfo, repInfo = headerEntry.split("_")[-2:]
+    rep = int(repInfo)
+    if genInfo == "Base":
+        gen = 0
+    else:
+        gen = int(genInfo.lstrip("F"))
+    return rep, gen
+
+
+def getFreqs(baseCounts):
+    aCount, tCount, cCount, gCount, nCount, delCount = [
+        int(x) for x in baseCounts.split(":")
+    ]
+    counts = [aCount, tCount, cCount, gCount, delCount]
+    denom = sum(counts)
+
+    return [x / denom for x in counts]
+
+
+def getWinningAlleleIndex(allFreqsForSnp, mode="final"):
+    assert allFreqsForSnp.shape == (7, 5)  # 7 timepoints X freqs for 5 possible alleles
+
+    freqScores = []
+    for alleleIndex in range(allFreqsForSnp.shape[1]):
+        if mode == "final":
+            freqScore = allFreqsForSnp[-1][alleleIndex]
+        elif mode == "velocity":
+            freqScore = allFreqsForSnp[-1][alleleIndex] - allFreqsForSnp[0][alleleIndex]
+        freqScores.append(freqScore)
+
+    return np.argmax(freqScores)
+
+
+def getMostCommonOtherAlleleIndex(allFreqsForSnp, winningAlleleIndex):
+    assert allFreqsForSnp.shape == (7, 5)  # 7 timepoints X freqs for 5 possible alleles
+
+    totalFreqs = []
+    for alleleIndex in range(allFreqsForSnp.shape[1]):
+        if alleleIndex == winningAlleleIndex:
+            totalFreqs.append(-1)
+        else:
+            totalFreq = sum(allFreqsForSnp[:, alleleIndex])
+            totalFreqs.append(totalFreq)
+
+    return np.argmax(totalFreqs)
+
+
+def encodeFreqsInPlace(freqArray):
+    for snpIndex in range(len(freqArray[0])):
+        allFreqsForSnp = []
+
+        for gen in range(len(freqArray)):
+            freqs = freqArray[gen][snpIndex]
+            allFreqsForSnp.append(freqs)
+
+        allFreqsForSnp = np.array(allFreqsForSnp)
+        winningAlleleIndex = getWinningAlleleIndex(allFreqsForSnp, mode="velocity")
+        otherAlleleIndex = getMostCommonOtherAlleleIndex(
+            allFreqsForSnp, winningAlleleIndex
+        )
+
+        for gen in range(len(freqArray)):
+            freqs = freqArray[gen][snpIndex]
+            winningFreq = freqs[winningAlleleIndex]
+            otherFreq = freqs[otherAlleleIndex]
+
+            freq = winningFreq / (winningFreq + otherFreq)
+            freqArray[gen][snpIndex] = freq
+        # assert winningFreq > 0 and winningFreq >= otherFreq
+
+def get_header(headerfile):
+    with open(headerfile, "r") as hfile:
+        header = hfile.readline().strip().split()
+        
+    return header
+
+def main(ua):
+    agp = argparse.ArgumentParser()
+    agp.add_argument("-i", "--infile", help="Sync file to convert to timesweeper NPZ format.")
+    agp.add_argument("-c", "--chrom", help="Chromosome to pull out of sync file.")
+    agp.add_argument("-o", "--outdir", help="Output directory.")
+    agp.add_argument("-w", "--window_size", help="Size of windows to extract.")
+    agp.add_argument("-h", "--headerfile", help="File containing a single line - the header of the syncfile. Must be provided if header is not present in the syncfile.")
+    ua = agp.parse_args()
+    
+    inFileName = ua.infile
+    outDir = ua.outdir 
+    winSize = ua.window_size 
+    targetChrom = ua.chrom 
+    
+    if ua.headerfile:
+       header = get_header(ua.headerfile)
+    else:
+        header = get_header(inFileName) 
+    
+    freqs = {}
+    for i in range(len(header)):
+        if "Dsim" in header[i]:
+            rep, gen = getRepAndGen(header[i])
+            if not rep in freqs:
+                freqs[rep] = {}
+            if not gen in freqs[rep]:
+                freqs[rep][gen] = {}
+
+
+    sys.stderr.write("reading snps and freqs\n")
+    positions = {}
+    with open(inFileName, "rt") as inFile:
+        for line in inFile:
+            line = line.strip().split()
+            chrom = line[0]
+            pos = int(line[1])
+            if chrom == targetChrom:
+                if not chrom in positions:
+                    positions[chrom] = []
+                    for rep in freqs:
+                        for gen in freqs[rep]:
+                            freqs[rep][gen][chrom] = []
+
+                for i in range(len(line)):
+                    if ":" in line[i]:
+                        rep, gen = getRepAndGen(header[i])
+                        currFreqs = getFreqs(line[i])
+                        freqs[rep][gen][chrom].append(currFreqs)
+                positions[chrom].append(pos)
+    sys.stderr.write("got all snps and freqs\n")
+
+
+    sys.stderr.write("formatting output\n")
+    for rep in freqs:
+        for chrom in positions:
+            outFileName = f"{outDir}/dsim_chrom_{chrom}_rep_{rep}.npz"
+
+            assert len(positions[chrom]) == len(freqs[rep][gen][chrom])
+
+            freqArray = []
+            for gen in sorted(freqs[rep]):
+                freqArray.append([])
+                for posIndex in range(len(positions[chrom])):
+                    currFreqs = freqs[rep][gen][chrom][posIndex]
+                    freqArray[-1].append(currFreqs)
+
+            encodeFreqsInPlace(freqArray)
+            freqArray = np.array(freqArray)
+            numGens = len(freqs[rep])
+            assert freqArray.shape == (numGens, len(positions[chrom]))
+
+            allFreqWins = []
+            allPosWins = []
+            startingIndices = range(len(positions[chrom]) - winSize)
+            for i in startingIndices:
+                currWin = freqArray[:, i : i + winSize]
+                allFreqWins.append(currWin)
+                currPositions = positions[chrom][i : i + winSize]
+                allPosWins.append(currPositions)
+            assert i + winSize == len(positions[chrom]) - 1
+
+            allFreqWins = np.array(allFreqWins)
+            allPosWins = np.array(allPosWins)
+            assert allFreqWins.shape == (len(startingIndices), numGens, winSize)
+            assert allPosWins.shape == (len(startingIndices), winSize)
+
+            np.savez(outFileName, aftIn=allFreqWins, aftInPosition=allPosWins)
+            
+    sys.stderr.write("all done!\n")
+
+    
+if __name__=="__main__":
+    main()
\ No newline at end of file
diff --git a/timesweeper/find_sweeps_npz.py b/timesweeper/find_sweeps_npz.py
index 5329363..6875f45 100644
--- a/timesweeper/find_sweeps_npz.py
+++ b/timesweeper/find_sweeps_npz.py
@@ -9,8 +9,7 @@
 from tensorflow.keras.models import load_model
 from tqdm import tqdm
 
-from utils.frequency_increment_test import fit
-from utils.gen_utils import write_fit, write_preds
+from utils.gen_utils import write_preds
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
@@ -62,29 +61,6 @@ def run_aft_windows(ts_aft, locs, chrom, model):
     return results_list
 
 
-def run_fit_windows(ts_aft, locs, chrom):
-    """
-    Iterates through windows of MAF time-series matrix and predicts using NN.
-    Args:
-        snps (list[tup(chrom, pos,  mut, s_coeff)]): Tuples of information for each SNP. Contains mut and s only if benchmarking == True.
-        genos (allel.GenotypeArray): Genotypes of all samples.
-        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
-        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
-        gens (list[int]): List of generations that were sampled.
-    Returns:
-        dict: P values from FIT.
-    """
-    gens = [i * 10 for i in range(7)]
-
-    results_list = []
-    for snp in tqdm(range(len(locs)), desc="Calculating FIT values"):
-        results_list.append(
-            (chrom, locs[snp, 25], fit(list(ts_aft[snp, :, 25]), gens)[1])
-        )  # tval, pval
-
-    return results_list
-
-
 def load_nn(model_path, summary=False):
     """
     Loads the trained Keras network.
@@ -101,25 +77,6 @@ def load_nn(model_path, summary=False):
     return model
 
 
-def write_fit(fit_list, outfile):
-    """
-    Writes FIT predictions to file.
-    Args:
-        fit_dict (dict): FIT p values and SNP information.
-        outfile (str): File to write results to.
-    """
-    chrom, bps, pval = zip(*fit_list)
-    inv_pval = [1 - i for i in pval]
-
-    predictions = pd.DataFrame({"Chrom": chrom, "BP": bps, "Inv_pval": inv_pval})
-    # predictions = predictions[predictions["Inv_pval"] > 0.9]
-
-    predictions.sort_values(["Chrom", "BP"], inplace=True)
-    predictions.to_csv(
-        os.path.join(outfile), header=True, index=False, sep="\t", float_format="%.3f"
-    )
-
-
 def write_preds(results_list, outfile, benchmark):
     """
     Writes NN predictions to file.
@@ -179,7 +136,3 @@ def main(ua):
     aft_predictions = run_aft_windows(ts_aft, locs, chrom, aft_model)
     write_preds(aft_predictions, f"{outdir}/aft_{chrom}_{rep}_preds.csv", ua.benchmark)
     logger.info(f"Done, results written to {outdir}/aft_{chrom}_{rep}_preds.csv")
-
-    # FIT
-    fit_predictions = run_fit_windows(ts_aft, locs, chrom)
-    write_fit(fit_predictions, f"{outdir}/fit_{chrom}_{rep}_preds.csv")
diff --git a/timesweeper/find_sweeps_vcf.py b/timesweeper/find_sweeps_vcf.py
index 5b5e719..a8194ad 100644
--- a/timesweeper/find_sweeps_vcf.py
+++ b/timesweeper/find_sweeps_vcf.py
@@ -1,68 +1,100 @@
-import logging
 import math
 import os
 
-import allel
 import numpy as np
+import pickle as pkl
+import pandas as pd
 from tensorflow.keras.models import load_model
 from tqdm import tqdm
 
-from .utils.frequency_increment_test import fit
-from .utils import snp_utils as su
-from .utils.gen_utils import read_config, write_fit, write_preds, get_logger
-from .utils import hap_utils as hu
+from timesweeper.make_training_features import prep_ts_aft, get_window_idxs
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import read_config, get_logger
+from timesweeper.utils import hap_utils as hu
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
 logger = get_logger("find_sweeps")
 
 
-def get_window_idxs(center_idx, win_size):
+def write_preds(results_dict, outfile, scaler, benchmark):
     """
-    Gets the win_size number of snps around a central snp.
+    Writes NN predictions to file.
 
     Args:
-        center_idx (int): Index of the central SNP to use for the window.
-        win_size (int): Size of window to use around the SNP, optimally odd number.
-
-    Returns:
-        list: Indices of all SNPs to grab for the feature matrix.
+        results_dict (dict): SNP NN prediction scores and window edges.
+        outfile (str): File to write results to.
     """
-    half_window = math.floor(win_size / 2)
-    return list(range(center_idx - half_window, center_idx + half_window + 1))
+    with open(scaler, "rb") as ifile:
+        scaler = pkl.load(ifile)
 
-
-def prep_ts_aft(genos, samp_sizes):
-    """
-    Iterates through timepoints and creates MAF feature matrices.
-
-    Args:
-        genos (allel.GenotypeArray): Genotype array containing all timepoints.
-        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
-
-    Returns:
-        np.arr: MAF array to use for predictions. Shape is (timepoints, MAF).
-    """
-    # Prep genos into time-series format and calculate Maft
-    ts_genos = su.split_arr(genos, samp_sizes)
-    min_alleles = su.get_vel_minor_alleles(ts_genos, np.max(genos))
-    ts_maft = []
-    for timepoint in ts_genos:
-        _genos = []
-        _genotypes = allel.GenotypeArray(timepoint).count_alleles(
-            max_allele=min_alleles.max()
+    lab_dict = {0: "Neut", 1: "SSV", 2: "SDN"}
+    if benchmark:
+        chroms, bps, mut_type, true_sel_coeff = zip(*results_dict.keys())
+    else:
+        chroms, bps = zip(*results_dict.keys())
+
+    neut_scores = [i[0][0] for i in results_dict.values()]
+    sdn_scores = [i[0][1] for i in results_dict.values()]
+    ssv_scores = [i[0][2] for i in results_dict.values()]
+    sdn_preds = [i[1] for i in results_dict.values()]
+    ssv_preds = [i[2] for i in results_dict.values()]
+    left_edges = [i[3] for i in results_dict.values()]
+    right_edges = [i[4] for i in results_dict.values()]
+    classes = [lab_dict[np.argmax(i[0])] for i in results_dict.values()]
+
+    sdn_s = scaler.inverse_transform(sdn_preds).squeeze().flatten()
+    ssv_s = scaler.inverse_transform(ssv_preds).squeeze().flatten()
+
+    if benchmark:
+        predictions = pd.DataFrame(
+            {
+                "Chrom": chroms,
+                "BP": bps,
+                "Mut_Type": mut_type,
+                "Class": classes,
+                "Sweep_Prob": [i + j for i, j in zip(sdn_scores, ssv_scores)],
+                "True_Sel_Coeff": true_sel_coeff,
+                "SDN_s_pred": sdn_s,
+                "SSV_s_pred": ssv_s,
+                "Neut_Prob": neut_scores,
+                "SDN_Prob": sdn_scores,
+                "SSV_Prob": ssv_scores,
+                "Win_Start": left_edges,
+                "Win_End": right_edges,
+            }
+        )
+    else:
+        predictions = pd.DataFrame(
+            {
+                "Chrom": chroms,
+                "BP": bps,
+                "Class": classes,
+                "Sweep_Prob": [i + j for i, j in zip(sdn_scores, ssv_scores)],
+                "SDN_s_pred": sdn_s,
+                "SSV_s_pred": ssv_s,
+                "Neut_Prob": neut_scores,
+                "SDN_Prob": sdn_scores,
+                "SSV_Prob": ssv_scores,
+                "Win_Start": left_edges,
+                "Win_End": right_edges,
+            }
         )
+        predictions = predictions[predictions["Neut_Prob"] < 0.5]
 
-        for snp, min_allele_idx in zip(_genotypes, min_alleles):
-            maf = su.calc_maft(snp, min_allele_idx)
-            _genos.append(maf)
+    predictions.sort_values(["Chrom", "BP"], inplace=True)
 
-        ts_maft.append(_genos)
+    predictions.to_csv(outfile, header=True, index=False, sep="\t")
 
-    return np.stack(ts_maft)
+    bed_df = predictions[["Chrom", "Win_Start", "Win_End", "BP"]]
+    bed_df.to_csv(
+        outfile.replace(".csv", ".bed"), mode="a", header=False, index=False, sep="\t"
+    )
 
 
-def run_aft_windows(snps, genos, samp_sizes, win_size, model):
+def run_aft_windows(
+    snps, genos, samp_sizes, win_size, class_model, sdn_model, ssv_model
+):
     """
     Iterates through windows of MAF time-series matrix and predicts using NN.
 
@@ -97,18 +129,28 @@ def run_aft_windows(snps, genos, samp_sizes, win_size, model):
         except Exception as e:
             logger.warning(f"Center {snps[center]} raised error {e}")
 
-    class_probs, sel_pred = model.predict(np.stack(data))
+    class_probs = class_model.predict(np.stack(data))
+    sdn_sel_preds = sdn_model.predict(np.stack(data))
+    ssv_sel_preds = ssv_model.predict(np.stack(data))
 
     results_dict = {}
-    for center, class_probs, sel_pred, l_e, r_e in zip(
-        centers, class_probs, sel_pred, left_edges, right_edges
+    for center, class_probs, sd, sv, l_e, r_e in zip(
+        centers, class_probs, sdn_sel_preds, ssv_sel_preds, left_edges, right_edges
     ):
-        results_dict[snps[center]] = (class_probs, sel_pred, l_e, r_e)
+        results_dict[snps[center]] = (
+            class_probs,
+            sd,
+            sv,
+            l_e,
+            r_e,
+        )
 
     return results_dict
 
 
-def run_hft_windows(snps, haps, ploidy, samp_sizes, win_size, model):
+def run_hft_windows(
+    snps, haps, ploidy, samp_sizes, win_size, class_model, sdn_model, ssv_model
+):
     """
     Iterates through windows of MAF time-series matrix and predicts using NN.
     Args:
@@ -127,71 +169,52 @@ def run_hft_windows(snps, haps, ploidy, samp_sizes, win_size, model):
     left_edges = []
     right_edges = []
     data = []
-    for center in tqdm(centers, desc="Predicting on HFS windows"):
-        try:
-            win_idxs = get_window_idxs(center, win_size)
-            window = np.swapaxes(haps[win_idxs, :], 0, 1)
-            str_window = hu.haps_to_strlist(window)
-            hfs = hu.getTSHapFreqs(str_window, [i * ploidy for i in samp_sizes])
-            data.append(hfs)
-            left_edges.append(snps[win_idxs[0]][1])
-            right_edges.append(snps[win_idxs[-1]][1])
-
-        except Exception as e:
-            logger.warning(f"Center {snps[center]} raised error {e}")
-            continue
-
-    probs = model.predict(np.stack(data))
+    for center in tqdm(centers, desc="Predicting on HFT windows"):
+        # try:
+        win_idxs = get_window_idxs(center, win_size)
+        window = np.swapaxes(haps[win_idxs, :], 0, 1)
+        str_window = hu.haps_to_strlist(window)
+        hft = hu.getTSHapFreqs(str_window, [i * ploidy for i in samp_sizes])
+        data.append(hft)
+        left_edges.append(snps[win_idxs[0]][1])
+        right_edges.append(snps[win_idxs[-1]][1])
+
+        # except Exception as e:
+        #    logger.warning(f"Center {snps[center]} raised error {e}")
+        #    continue
+
+    class_probs = class_model.predict(np.stack(data))
+    sdn_sel_preds = sdn_model.predict(np.stack(data))
+    ssv_sel_preds = ssv_model.predict(np.stack(data))
 
     results_dict = {}
-    for center, prob, l_e, r_e in zip(centers, probs, left_edges, right_edges):
-        results_dict[snps[center]] = (prob, l_e, r_e)
+    for center, class_probs, sd, sv, l_e, r_e in zip(
+        centers, class_probs, sdn_sel_preds, ssv_sel_preds, left_edges, right_edges
+    ):
+        results_dict[snps[center]] = (
+            class_probs,
+            sd,
+            sv,
+            l_e,
+            r_e,
+        )
 
     return results_dict
 
 
-def run_fit_windows(snps, genos, samp_sizes, win_size, gens):
+def get_window_idxs(center_idx, win_size):
     """
-    Iterates through windows of MAF time-series matrix and predicts using NN.
+    Gets the win_size number of snps around a central snp.
 
     Args:
-        snps (list[tup(chrom, pos,  mut)]): Tuples of information for each SNP. Contains mut only if benchmarking == True.
-        genos (allel.GenotypeArray): Genotypes of all samples.
-        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
-        win_size (int): Number of SNPs to use for each prediction. Needs to match how NN was trained.
-        gens (list[int]): List of generations that were sampled.
+        center_idx (int): Index of the central SNP to use for the window.
+        win_size (int): Size of window to use around the SNP, optimally odd number.
 
     Returns:
-        dict: P values from FIT.
+        list: Indices of all SNPs to grab for the feature matrix.
     """
-    ts_aft = prep_ts_aft(genos, samp_sizes)
-    results_dict = {}
-    buffer = int(win_size / 2)
-    for idx in tqdm(range(buffer, len(snps) - buffer), desc="Calculating FIT values"):
-        results_dict[snps[idx]] = fit(list(ts_aft[:, idx]), gens)  # tval, pval
-
-    return results_dict
-
-
-def run_fet_windows(genos, samp_sizes):
-    ts_genos = su.split_arr(genos, samp_sizes)
-    min_alleles = su.get_vel_minor_alleles(ts_genos, np.max(genos))
-
-    # get count values for maj, min using snp utils
-    # feed into fet
-    # collect fet values into fet_geno counts for all snps
-    # fet_pvals = list with len(snps)
-
-    for timepoint in [0, -1]:
-        _genotypes = allel.GenotypeArray(timepoint).count_alleles(
-            max_allele=min_alleles.max()
-        )
-        min_allele_counts = []
-        fet_geno_counts = []
-        for snp, min_allele_idx in zip(_genotypes, min_alleles):
-            fet_geno_counts.append(min_allele_counts)
-
-    raise NotImplementedError
+    half_window = math.floor(win_size / 2)
+    return list(range(center_idx - half_window, center_idx + half_window + 1))
 
 
 def load_nn(model_path, summary=False):
@@ -214,24 +237,28 @@ def load_nn(model_path, summary=False):
 
 def main(ua):
     yaml_data = read_config(ua.yaml_file)
-    (work_dir, samp_sizes, ploidy, win_size, outdir, aft_model) = (
-        yaml_data["work dir"],
-        yaml_data["sample sizes"],
-        yaml_data["ploidy"],
-        yaml_data["win_size"],
-        ua.outdir,
-        load_nn(ua.aft_model),
-    )
-
-    # If you're doing simple sims you probably aren't calculating years out
-    if "gens sampled" in yaml_data:
-        gens_sampled = yaml_data["gens sampled"]
-    else:
-        gens_sampled = None
-
-    outdir = ua.outdir
-    if not os.path.exists(outdir):
-        os.makedirs(outdir)
+    samp_sizes = yaml_data["sample sizes"]
+    ploidy = yaml_data["ploidy"]
+    win_size = yaml_data["win_size"]
+
+    class_aft_model = load_nn(ua.aft_class_model)
+    if "ssv" in ua.aft_reg_model:
+        aft_ssv_model = load_model(ua.aft_reg_model)
+        aft_sdn_model = load_model(str(ua.aft_reg_model).replace("ssv", "sdn"))
+    elif "sdn" in ua.aft_reg_model:
+        aft_ssv_model = load_model(str(ua.aft_reg_model).replace("ssv", "sdn"))
+        aft_sdn_model = load_model(ua.aft_reg_model)
+
+    if ua.hft_class_model:
+        class_hft_model = load_nn(ua.hft_class_model)
+        if "ssv" in ua.hft_reg_model:
+            hft_ssv_model = load_model(ua.hft_reg_model)
+            hft_sdn_model = load_model(str(ua.hft_reg_model).replace("ssv", "sdn"))
+        elif "sdn" in ua.hft_reg_model:
+            hft_ssv_model = load_model(str(ua.hft_reg_model).replace("ssv", "sdn"))
+            hft_sdn_model = load_model(ua.hft_reg_model)
+
+    outfile = ua.outfile
 
     # Chunk and iterate for NN predictions to not take up too much space
     vcf_iter = su.get_vcf_iter(ua.input_vcf, ua.benchmark)
@@ -240,42 +267,37 @@ def main(ua):
         logger.info(f"Processing VCF chunk {chunk_idx}")
 
         # aft
-        try:
-            genos, snps = su.vcf_to_genos(chunk, ua.benchmark)
-            aft_predictions = run_aft_windows(
-                snps, genos, samp_sizes, win_size, aft_model
-            )
-            write_preds(aft_predictions, f"{outdir}/aft_preds.csv", ua.benchmark)
+        # try:
+        genos, snps = su.vcf_to_genos(chunk, ua.benchmark)
+        aft_predictions = run_aft_windows(
+            snps,
+            genos,
+            samp_sizes,
+            win_size,
+            class_aft_model,
+            aft_sdn_model,
+            aft_ssv_model,
+        )
+        write_preds(aft_predictions, f"{outfile}_aft.csv", ua.scalar, ua.benchmark)
 
-        except Exception as e:
-            logger.error(f"Cannot process chunk {chunk_idx} using AFT due to {e}")
+        # except Exception as e:
+        #    logger.error(f"Cannot process chunk {chunk_idx} using AFT due to {e}")
 
         # hft
-        if ua.hft_model:
-            try:
-                haps, snps = su.vcf_to_haps(chunk, ua.benchmark)
-                aft_predictions = run_hft_windows(
-                    snps, haps, ploidy, samp_sizes, win_size, load_nn(ua.hft_model)
-                )
-                write_preds(aft_predictions, f"{outdir}/hft_preds.csv", ua.benchmark)
-
-            except Exception as e:
-                logger.error(f"Cannot process chunk {chunk_idx} using HFT due to {e}")
-
-    if gens_sampled:
-        vcf_iter = su.get_vcf_iter(ua.input_vcf, ua.benchmark)
-        for chunk_idx, chunk in enumerate(vcf_iter):
-            chunk = chunk[0]  # Why you gotta do me like that, skallel?
-
-            # FIT
+        if ua.hft_class_model:
             # try:
-            genos, snps = su.vcf_to_genos(chunk, ua.benchmark)
-            fit_predictions = run_fit_windows(
-                snps, genos, samp_sizes, win_size, gens_sampled
+            haps, snps = su.vcf_to_haps(chunk, ua.benchmark)
+            hft_predictions = run_hft_windows(
+                snps,
+                haps,
+                ploidy,
+                samp_sizes,
+                win_size,
+                class_hft_model,
+                hft_sdn_model,
+                hft_ssv_model,
             )
-            write_fit(fit_predictions, f"{outdir}/fit_preds.csv", ua.benchmark)
-            # except Exception as e:
-            #    logger.error(f"Cannot process chunk {chunk_idx} using AFT due to {e}")
+            write_preds(hft_predictions, f"{outfile}_hft.csv", ua.scalar, ua.benchmark)
 
-    else:
-        logger.info("Cannot calculate FIT, years sampled and gen time not supplied.")
+            # except Exception as e:
+            #    logger.error(f"Cannot process chunk {chunk_idx} using HFT due to {e}")
diff --git a/timesweeper/make_training_features.py b/timesweeper/make_training_features.py
index 7ef1daf..9baf371 100644
--- a/timesweeper/make_training_features.py
+++ b/timesweeper/make_training_features.py
@@ -1,19 +1,23 @@
 import logging
+import math
 import multiprocessing as mp
 import os
 import pickle
 import sys
 from glob import glob
 from itertools import cycle
-import random
+from random import sample
+import allel
 import numpy as np
 from numpy.random import default_rng
 from tqdm import tqdm
-
-from .find_sweeps_vcf import prep_ts_aft, get_window_idxs
-from .utils import snp_utils as su
-from .utils.gen_utils import get_rep_id, read_config, get_scenario_from_filename
-from .utils.hap_utils import haps_to_strlist, getTSHapFreqs
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import (
+    get_rep_id,
+    get_scenario_from_filename,
+    read_config,
+)
+from timesweeper.utils.hap_utils import getTSHapFreqs, haps_to_strlist
 
 logging.basicConfig()
 logger = logging.getLogger("make_training_feats")
@@ -24,6 +28,11 @@
 warnings.filterwarnings("error")
 
 
+def draw_rand_center_offset(max_offset=700):
+    rng = default_rng(np.random.seed(int.from_bytes(os.urandom(4), byteorder="little")))
+    return int(rng.uniform(-max_offset, max_offset, 1)[0])
+
+
 def add_missingness(data, m_rate, nan_val=-1):
     rng = default_rng(np.random.seed(int.from_bytes(os.urandom(4), byteorder="little")))
     missing = rng.binomial(1, m_rate, data.shape)
@@ -32,7 +41,105 @@ def add_missingness(data, m_rate, nan_val=-1):
     return data
 
 
-def get_aft_central_window(snps, genos, samp_sizes, win_size, missingness, mut_types):
+def check_freq_increase(ts_afs, min_increase=0.25):
+    """Quick test to make sure a given example is increasing properly, can be used to filter training."""
+    center = int(ts_afs.shape[0] / 2)
+    if ts_afs[-1, center] - ts_afs[0, center] >= min_increase:
+        return True
+    else:
+        return False
+
+
+def subsample_tps(inds_per_tp, og_tps, num_sampled_tps):
+    """Subsamples timepoints during VCF by sampling samples list."""
+    samples_list = []
+    if num_sampled_tps == 1:
+        t = og_tps + 1  # VCFs merged samples are 1-indexed, get last
+        samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+    elif num_sampled_tps == 2:
+        # This covers the first sampling point
+        samples_list.extend([f"i{i}" for i in range(inds_per_tp)])
+        t = og_tps + 1  # VCFs merged samples are 1-indexed, get last
+        samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+    else:
+        tps = np.linspace(1, og_tps, num_sampled_tps, dtype=int)
+        # This covers the first sampling point
+        samples_list.extend([f"i{i}" for i in range(inds_per_tp)])
+        # This gets the rest
+        for t in tps[1:]:
+            samples_list.extend([f"{t}:i{i}" for i in range(inds_per_tp)])
+
+    return samples_list
+
+
+def subsample_inds(inds_per_tp, subsample_size, num_tps):
+    """Subsamples individuals from each timepoint to re-use simulations."""
+    samples_list = []
+    samples_list.extend(
+        [f"i{i}" for i in np.random.choice(inds_per_tp, subsample_size, replace=False)]
+    )
+    for s in range(2, num_tps + 1):
+        samples_list.extend(
+            [
+                f"{s}:i{i}"
+                for i in np.random.choice(inds_per_tp, subsample_size, replace=False)
+            ]
+        )
+    return samples_list
+
+
+def get_window_idxs(center_idx, win_size):
+    """
+    Gets the win_size number of snps around a central snp.
+
+    Args:
+        center_idx (int): Index of the central SNP to use for the window.
+        win_size (int): Size of window to use around the SNP, optimally odd number.
+
+    Returns:
+        list: Indices of all SNPs to grab for the feature matrix.
+    """
+    half_window = math.floor(win_size / 2)
+    return list(range(center_idx - half_window, center_idx + half_window + 1))
+
+
+def prep_ts_aft(genos, samp_sizes):
+    """
+    Iterates through timepoints and creates MAF feature matrices.
+
+    Args:
+        genos (allel.GenotypeArray): Genotype array containing all timepoints.
+        samp_sizes (list[int]): Number of chromosomes sampled at each timepoint.
+
+    Returns:
+        np.arr: MAF array to use for predictions. Shape is (timepoints, MAF).
+    """
+    # Prep genos into time-series format and calculate Maft
+    ts_genos = su.split_arr(genos, samp_sizes)
+
+    min_alleles, first_genos, last_genos = su.get_vel_minor_alleles(
+        ts_genos, np.max(genos)
+    )
+
+    ts_maft = []
+    for timepoint in ts_genos:
+        _genos = []
+        _genotypes = allel.GenotypeArray(timepoint).count_alleles(
+            max_allele=min_alleles.max()
+        )
+
+        for snp, min_allele_idx in zip(_genotypes, min_alleles):
+            maf = su.calc_maft(snp, min_allele_idx)
+            _genos.append(maf)
+
+        ts_maft.append(_genos)
+
+    return np.stack(ts_maft)
+
+
+def get_aft_central_window(
+    snps, genos, samp_sizes, win_size, missingness, mut_types, offset
+):
     """
     Check for whether a non-control mutation type is present. If not, return the central-most mutation.
 
@@ -54,23 +161,44 @@ def get_aft_central_window(snps, genos, samp_sizes, win_size, missingness, mut_t
     centers = range(buffer, len(snps) - buffer)
 
     center_idx = int(len(centers) / 2)
-    for center in centers:
-        if snps[center][2] in mut_types:
-            center_idx = center
-            break
+
+    try:
+        for center in centers:
+            if snps[center][2] in mut_types:
+                center_idx = center
+                break
+            else:
+                pass
+    except:
+        pass
+
+    sel_coeff = snps[center_idx][3]
+
+    if offset:
+        rng = default_rng(
+            np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+        )
+        if rng.uniform(0, 3, 1)[0] > 2:
+            rand_offset = draw_rand_center_offset()
+            center_idx += rand_offset
         else:
-            pass
+            rand_offset = 0
+    else:
+        rand_offset = 0
 
     win_idxs = get_window_idxs(center_idx, win_size)
+
     window = ts_aft[:, win_idxs]
     center_aft = window
-    sel_coeff = snps[center_idx][3]
-    missing_center_aft = add_missingness(center_aft, m_rate=missingness)
 
-    return missing_center_aft, sel_coeff
+    missing_center_aft = add_missingness(
+        center_aft, m_rate=missingness
+    )  # If no missingness, will just return
+
+    return missing_center_aft, sel_coeff, rand_offset
 
 
-def get_hft_central_window(snps, haps, samp_sizes, win_size, mut_types):
+def get_hft_central_window(snps, haps, samp_sizes, win_size, mut_types, offset):
     """
     Iterates through windows of MAF time-series matrix and gets the central window.
     Does not have as many utility functions as AFT such as missingness and variable sorting methods.
@@ -87,46 +215,67 @@ def get_hft_central_window(snps, haps, samp_sizes, win_size, mut_types):
     centers = range(buffer, len(snps) - buffer)
 
     center_idx = int(len(centers) / 2)
-    for center in centers:
-        if snps[center][2] in mut_types:
-            center_idx = center
-            break
+
+    try:
+        for center in centers:
+            if snps[center][2] in mut_types:
+                center_idx = center
+                break
+            else:
+                pass
+    except:
+        pass
+
+    sel_coeff = snps[center_idx][3]
+
+    if offset:
+        rng = default_rng(
+            np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
+        )
+        if rng.uniform(1, 3, 1)[0] > 2:
+            rand_offset = draw_rand_center_offset()
+            center_idx += rand_offset
+
         else:
-            pass
+            rand_offset = 0
+    else:
+        rand_offset = 0
 
     win_idxs = get_window_idxs(center_idx, win_size)
     window = np.swapaxes(haps[win_idxs, :], 0, 1)
     str_window = haps_to_strlist(window)
     central_hfs = getTSHapFreqs(str_window, samp_sizes)
-    sel_coeff = snps[center_idx][3]
 
-    return central_hfs, sel_coeff
-
-
-def check_freq_increase(ts_afs, min_increase=0.25):
-    """Quick test to make sure a given example is increasing properly, can be used to filter training."""
-    center = int(ts_afs.shape[0] / 2)
-    if ts_afs[-1, center] - ts_afs[0, center] >= min_increase:
-        return True
-    else:
-        return False
+    return central_hfs, sel_coeff, rand_offset
 
 
 def aft_worker(
-    in_vcf, mut_types, scenarios, samp_sizes, win_size, missingness, verbose=False,
+    in_vcf,
+    mut_types,
+    scenarios,
+    samp_sizes,
+    samps_list,
+    win_size,
+    offset,
+    missingness,
+    verbose=False,
 ):
-    benchmark = True
+    benchmark = True  # Want to get all the info we can from sims in training
     try:
         id = get_rep_id(in_vcf)
         scenario = get_scenario_from_filename(in_vcf, scenarios)
-        vcf = su.read_vcf(in_vcf, benchmark)
+
+        vcf = su.read_vcf(in_vcf, samps_list, benchmark)
         genos, snps = su.vcf_to_genos(vcf, benchmark)
 
-        central_aft, sel_coeff = get_aft_central_window(
-            snps, genos, samp_sizes, win_size, missingness, mut_types
+        central_aft, sel_coeff, rand_offset = get_aft_central_window(
+            snps, genos, samp_sizes, win_size, missingness, mut_types, offset
         )
 
-        return id, scenario, central_aft, sel_coeff
+        if "neut" not in scenario.lower() and sel_coeff == 0.0:
+            raise Exception
+
+        return id, scenario, central_aft, sel_coeff, rand_offset
 
     except UserWarning as Ue:
         print(Ue)
@@ -142,20 +291,29 @@ def aft_worker(
 
 
 def hft_worker(
-    in_vcf, mut_types, scenarios, samp_sizes, win_size, ploidy=1, verbose=False,
+    in_vcf,
+    mut_types,
+    scenarios,
+    samp_sizes,
+    samps_list,
+    win_size,
+    offset,
+    ploidy=2,
+    verbose=False,
 ):
     benchmark = True
     try:
         id = get_rep_id(in_vcf)
         scenario = get_scenario_from_filename(in_vcf, scenarios)
-        vcf = su.read_vcf(in_vcf, benchmark)
+
+        vcf = su.read_vcf(in_vcf, samps_list, benchmark)
         haps, snps = su.vcf_to_haps(vcf, benchmark)
 
-        central_hft, sel_coeff = get_hft_central_window(
-            snps, haps, [ploidy * i for i in samp_sizes], win_size, mut_types,
+        central_hft, sel_coeff, rand_offset = get_hft_central_window(
+            snps, haps, [ploidy * i for i in samp_sizes], win_size, mut_types, offset
         )
 
-        return id, scenario, central_hft, sel_coeff
+        return id, scenario, central_hft, sel_coeff, rand_offset
 
     except UserWarning as Ue:
         # print(Ue)
@@ -181,15 +339,37 @@ def main(ua):
         yaml_data["win_size"],
         ua.threads,
     )
+    if ua.subsample_inds:
+        if ua.subsample_tps:
+            logger.error("Can't subsample both timepoints and individuals.")
 
-    filelist = glob(f"{work_dir}/vcfs/*/*/merged.vcf", recursive=True)[:10]
+        samps_list = subsample_inds(samp_sizes[0], ua.subsample_inds, len(samp_sizes))
+    else:
+        samps_list = None
+
+    if ua.subsample_tps:
+        if not ua.og_tps:
+            logger.error(
+                "Must provide original simulation number of tps in order to properly subset."
+            )
+
+        samps_list = subsample_tps(samp_sizes[0], ua.og_tps, ua.subsample_tps)
+
+    if ua.allow_shoulders:
+        offset = int(ua.allow_shoulders)
+    else:
+        offset = 0
+
+    filelist = glob(f"{work_dir}/vcfs/*/*/merged.vcf", recursive=True)
 
     aft_work_args = zip(
         filelist,
         cycle([mut_types]),
         cycle([scenarios]),
         cycle([samp_sizes]),
+        cycle([samps_list]),
         cycle([win_size]),
+        cycle([offset]),
         cycle([ua.missingness]),
         cycle([ua.verbose]),
     )
@@ -198,60 +378,70 @@ def main(ua):
         cycle([mut_types]),
         cycle([scenarios]),
         cycle([samp_sizes]),
+        cycle([samps_list]),
         cycle([win_size]),
+        cycle([offset]),
         cycle([int(ploidy)]),
         cycle([ua.verbose]),
     )
 
-    pool = mp.Pool(threads)
-    if ua.no_progress:
-        aft_work_res = pool.starmap(aft_worker, aft_work_args, chunksize=4,)
+    debug = False
+    if debug:
+        aft_work_res = [aft_worker(*i) for i in aft_work_args]
+        hft_work_res = [hft_worker(*i) for i in hft_work_args]
+    else:
+        pool = mp.Pool(threads)
+        if ua.no_progress:
+            aft_work_res = pool.starmap(aft_worker, aft_work_args, chunksize=4,)
 
-        if ua.hft:
-            hft_work_res = pool.starmap(hft_worker, hft_work_args, chunksize=4,)
+            if ua.hft:
+                hft_work_res = pool.starmap(hft_worker, hft_work_args, chunksize=4,)
 
-        pool.close()
-    else:
-        aft_work_res = pool.starmap(
-            aft_worker,
-            tqdm(
-                aft_work_args, desc="Formatting AFT training data", total=len(filelist),
-            ),
-            chunksize=4,
-        )
-        if ua.hft:
-            hft_work_res = pool.starmap(
-                hft_worker,
+            pool.close()
+        else:
+            aft_work_res = pool.starmap(
+                aft_worker,
                 tqdm(
-                    hft_work_args,
-                    desc="Formatting HFT training data",
+                    aft_work_args,
+                    desc="Formatting AFT training data",
                     total=len(filelist),
                 ),
                 chunksize=4,
             )
-        pool.close()
+            if ua.hft:
+                hft_work_res = pool.starmap(
+                    hft_worker,
+                    tqdm(
+                        hft_work_args,
+                        desc="Formatting HFT training data",
+                        total=len(filelist),
+                    ),
+                    chunksize=4,
+                )
+            pool.close()
 
-    sys.exit(0)
     pickle_dict = {}
     for s in scenarios:
         pickle_dict[s] = {}
 
     for res in aft_work_res:
         if res:
-            rep, scenario, aft, s = res
+            rep, scenario, aft, s, off = res
 
             pickle_dict[scenario][rep] = {}
             pickle_dict[scenario][rep]["aft"] = aft
             pickle_dict[scenario][rep]["sel_coeff"] = s
+            pickle_dict[scenario][rep]["center_offset"] = off
 
     if ua.hft:
         for res in hft_work_res:
             try:
                 if res:
-                    rep, scenario, hft, s = res
+                    rep, scenario, hft, s, off = res
                     if rep in pickle_dict[scenario].keys():
                         pickle_dict[scenario][rep]["hft"] = hft
                         pickle_dict[scenario][rep]["sel_coeff"] = s
+                        pickle_dict[scenario][rep]["center_offset"] = off
 
             except KeyError as e:
                 print(e)
diff --git a/timesweeper/merge_logs_acc.py b/timesweeper/merge_logs_acc.py
new file mode 100644
index 0000000..b5ff396
--- /dev/null
+++ b/timesweeper/merge_logs_acc.py
@@ -0,0 +1,66 @@
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+
+def main(ua):
+    if not os.path.exists(ua.output_dir):
+        os.makedirs(ua.output_dir)
+
+    log_df = pd.read_csv(ua.summary_tsv, header=0, sep="\t")
+    acc_df = pd.read_csv(ua.test_file, header=0, sep=",")
+    merged = pd.merge(
+        right=acc_df,
+        left=log_df,
+        how="right",
+        left_on=["rep", "sweep"],
+        right_on=["rep", "class"],
+    )
+    merged["abs_error"] = np.abs(merged["pred_sel_coeff"] - merged["true_sel_coeff"])
+    merged.to_csv(
+        f"{ua.output_dir}/merged_table.tsv", index=False, header=True, sep="\t"
+    )
+
+    merged.plot(
+        x="selCoeff",
+        y="abs_error",
+        kind="scatter",
+        title="Absolute S Prediction Error vs Selection Coefficient",
+        xlabel="Selection Coefficient s",
+        ylabel="Absolute Error | abs(predicted s - true s)",
+        xlim=(0.0, 0.25),
+        ylim=(0.0, 0.25),
+    )
+    plt.savefig(f"{ua.output_dir}/error_vs_selcoeff.png")
+    plt.clf()
+
+    merged.plot(
+        x="sampOffset",
+        y="abs_error",
+        kind="scatter",
+        title="Absolute S Prediction Error vs Selection Coefficient",
+        xlabel="Sampling Offset from Selection",
+        ylabel="Absolute Error | abs(predicted s - true s)",
+        xlim=(-50, 50),
+        ylim=(0.0, 0.25),
+        xticks=range(-50, 60, 10),
+    )
+    plt.savefig(f"{ua.output_dir}/error_vs_sampletime.png")
+    plt.clf()
+
+    merged.plot(
+        x="sampOffset",
+        y="selCoeff",
+        kind="scatter",
+        title="Sampling Offset from Selection vs Selection Coefficient",
+        xlabel="Sampling Offset from Selection",
+        ylabel="Selection Coefficient",
+        xlim=(-50, 50),
+        ylim=(0.0, 0.25),
+        xticks=range(-50, 60, 10),
+    )
+    plt.savefig(f"{ua.output_dir}/selcoeff_vs_sampletime.png")
+    plt.clf()
diff --git a/timesweeper/models.py b/timesweeper/models.py
index 316535a..8d26f2a 100644
--- a/timesweeper/models.py
+++ b/timesweeper/models.py
@@ -1,11 +1,9 @@
-from keras import layers
-from keras.models import Model
+from tensorflow.keras import layers
+from tensorflow.keras.models import Model
 
 # fmt: off
 def create_TS_class_model(datadim, n_class):
     """
-    Creates Time-Distributed SHIC model that uses 3 convlutional blocks with concatenation.
-
     Returns:
         Model: Keras compiled model.
     """
@@ -16,7 +14,7 @@ def create_TS_class_model(datadim, n_class):
     h = layers.Dropout(0.15)(h)
     h = layers.Flatten()(h)
 
-    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dense(512, activation="relu")(h)
     h = layers.Dropout(0.2)(h)        
     h = layers.Dense(264, activation="relu")(h)
     h = layers.Dropout(0.2)(h)
@@ -35,8 +33,6 @@ def create_TS_class_model(datadim, n_class):
 
 def create_TS_reg_model(datadim):
     """
-    Creates Time-Distributed SHIC model that uses 3 convlutional blocks with concatenation.
-
     Returns:
         Model: Keras compiled model.
     """
@@ -53,7 +49,7 @@ def create_TS_reg_model(datadim):
     h = layers.Dropout(0.2)(h)
     h = layers.Dense(128, activation="relu")(h)
     h = layers.Dropout(0.1)(h)
-    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+    reg_output = layers.Dense(1, activation="linear", name="reg_output")(h)
 
     model = Model(inputs=[model_in], outputs=[reg_output], name="Timesweeper_Reg")
     model.compile(
@@ -64,9 +60,212 @@ def create_TS_reg_model(datadim):
 
     return model
 
+
+# fmt: off
+def create_2D_TS_class_model(datadim, n_class):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Conv2D(64, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv2D(64, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling2D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(512, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="2DTimesweeper_Class")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+def create_2D_TS_reg_model(datadim):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    
+    h = layers.Conv2D(64, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv2D(64, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling2D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(264, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="2DTimesweeper_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
+
+
+# fmt: off
+def create_big_TS_class_model(datadim, n_class):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Conv1D(1028, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv1D(1028, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling1D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(1028, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(512, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="Big_Timesweeper_Class")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+def create_big_TS_reg_model(datadim):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Conv1D(1028, 3, activation="relu", padding="same")(model_in)
+    h = layers.Conv1D(1028, 3, activation="relu", padding="same")(h)
+    h = layers.MaxPooling1D(pool_size=3, padding="same")(h)
+    h = layers.Dropout(0.15)(h)
+    h = layers.Flatten()(h)
+
+    h = layers.Dense(1028, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(512, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="Big_Timesweeper_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
+
+
+# fmt: off
+def create_1tp_class_model(datadim, n_class):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+
+    h = layers.Dense(264, activation="relu")(model_in)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(56, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="Timesweeper_Class")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+
+def create_1tp_reg_model(datadim):
+    """
+    Returns:
+        Model: Keras compiled model.
+    """
+    model_in = layers.Input(datadim)
+    h = layers.Dense(264, activation="relu")(model_in)
+    h = layers.Dropout(0.2)(h)        
+    h = layers.Dense(128, activation="relu")(h)
+    h = layers.Dropout(0.2)(h)
+    h = layers.Dense(56, activation="relu")(h)
+    h = layers.Dropout(0.1)(h)
+
+    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="Timesweeper_Reg")
+    model.compile(
+        loss={"reg_output":"mse"},
+        optimizer="adam",
+        metrics={"reg_output": "mse"},
+    )
+
+    return model
 # fmt: on
 
 
+def create_rnn_class_model(datadim, n_class):
+    model_in = layers.Input(datadim)
+    h = layers.LSTM(100, input_shape=datadim, activation="tanh")(model_in)
+    h = layers.Dropout(0.1)(h)
+
+    class_output = layers.Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output], name="Timesweeper_Class")
+    model.compile(
+        loss={"class_output": "categorical_crossentropy"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy"},
+    )
+
+    return model
+
+
+def create_rnn_reg_model(datadim):
+    model_in = layers.Input(datadim)
+    h = layers.LSTM(100, input_shape=datadim, activation="tanh")(model_in)
+    h = layers.Dropout(0.1)(h)
+
+    reg_output = layers.Dense(1, activation="relu", name="reg_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[reg_output], name="RNN_Reg")
+    model.compile(
+        loss={"reg_output": "mse"}, optimizer="adam", metrics={"reg_output": "mse"},
+    )
+
+    return model
+
+
 def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
     # Attention and Normalization
     x = layers.MultiHeadAttention(
@@ -100,7 +299,7 @@ def create_transformer_reg_model(
     for _ in range(num_transformer_blocks):
         x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
 
-    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
+    x = layers.GlobalAveragePooling1D(data_format="channels_last")(x)
     for dim in mlp_units:
         x = layers.Dense(dim, activation="relu")(x)
         x = layers.Dropout(mlp_dropout)(x)
@@ -131,7 +330,7 @@ def create_transformer_class_model(
     for _ in range(num_transformer_blocks):
         x = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
 
-    x = layers.GlobalAveragePooling1D(data_format="channels_first")(x)
+    x = layers.GlobalAveragePooling1D(data_format="channels_last")(x)
     for dim in mlp_units:
         x = layers.Dense(dim, activation="relu")(x)
         x = layers.Dropout(mlp_dropout)(x)
diff --git a/timesweeper/net_hyper_tuning.py b/timesweeper/net_hyper_tuning.py
new file mode 100644
index 0000000..16304eb
--- /dev/null
+++ b/timesweeper/net_hyper_tuning.py
@@ -0,0 +1,308 @@
+import logging
+import os
+import pickle
+import random
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from sklearn.metrics import confusion_matrix, mean_absolute_error
+from sklearn.model_selection import train_test_split
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.layers import Conv1D, Dense, Dropout, Flatten, Input, MaxPooling1D
+from tensorflow.keras.models import Model, save_model
+from tensorflow.keras.utils import to_categorical
+from tensorflow.keras.utils import plot_model
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+from sklearn.model_selection import GridSearchCV
+from tensorflow.keras.wrappers.scikit_learn import KerasClassifier
+
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = 42
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+
+
+
+grid = ParameterGrid([{'kernel_size': [16, 32, 64, 128, 256]},
+                      {'kernel': ['rbf'], 'gamma': [1, 10]}]
+
+list(ParameterGrid(grid)) == [{'kernel': 'linear'},
+                              {'kernel': 'rbf', 'gamma': 1},
+                              {'kernel': 'rbf', 'gamma': 10}]
+
+
+
+def get_mlp_model(datadim, n_class,
+                  kernel_sizeLayerer_a=64,
+                  kernel_sizeLayerer_b=64,
+	              dropoutCovLayer_a=0.15,
+                  learnRate=0.01,
+                  add_one_layer=False,
+                  add_two_layer=False,
+                  dropoutFullLayer_a=0.2,
+                  dropoutFullLayer_b=0.2,
+                  dropoutFullLayer_c=0.1):
+    """
+        Create a CNN model based on the hypter parameter input for the whole sample
+
+    Args:
+          kernel_sizeLayerer_a (int): the kernel size of the first convolutiona layer in the first convlutional block.
+          kernel_sizeLayerer_b (int): the kernel size of the second convolutiona layer in the first convlutional block.
+          dropoutCovLayer_a (float): the dropout rate in the first convlutional block.
+          learnRate (float): the learning rate of the CNN.
+          add_one_layer (boolean): if true, add one additional convlutional block.
+          add_two_layer (boolean): if true, add a second additional convlutional block.
+          dropoutFullLayer_a (float): the dropout rate of the first layer in the fully connnected block.
+          dropoutFullLayer_b (float): the dropout rate of the second layer in the fully connnected block.
+          dropoutFullLayer_c (float): the dropout rate of the third layer in the fully connnected block.
+
+    Returns:
+        kera.classifier: a cnn fitted to the input hyperparameter.
+    """
+    ## Convolution Layer One
+    model_in = Input(datadim)
+    h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(model_in)
+    h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+    h = MaxPooling1D(pool_size=3, padding="same")(h)
+    h = Dropout(dropoutCovLayer_a)(h)
+    h = Flatten()(h)
+
+    ## Convolution Layer Two if add_one_layer
+    if add_one_layer:
+        h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(h)
+        h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+        h = MaxPooling1D(pool_size=3, padding="same")(h)
+        h = Dropout(dropoutCovLayer_a)(h)
+        h = Flatten()(h)
+
+    ## Convolution Layer Three if add_two_layer
+    if add_two_layer:
+        h = Conv1D(kernel_sizeLayerer_a, 3, activation="relu", padding="same")(h)
+        h = Conv1D(kernel_sizeLayerer_b, 3, activation="relu", padding="same")(h)
+        h = MaxPooling1D(pool_size=3, padding="same")(h)
+        h = Dropout(dropoutCovLayer_a)(h)
+        h = Flatten()(h)
+
+    ## Fully Connected Layer
+    h = Dense(264, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_a)(h)
+    h = Dense(264, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_b)(h)
+    h = Dense(128, activation="relu")(h)
+    h = Dropout(dropoutFullLayer_c)(h)
+
+    reg_output = Dense(1, activation="relu", name="reg_output")(h)
+    class_output = Dense(n_class, activation="softmax", name="class_output")(h)
+
+    model = Model(inputs=[model_in], outputs=[class_output, reg_output], name="Timesweeper")
+    model.compile(
+        loss={"class_output":"categorical_crossentropy", "reg_output":"mse"},
+        optimizer="adam",
+        metrics={"class_output": "accuracy", "reg_output": "mse"},
+    )
+
+    return model
+
+
+
+def main(ua):
+    if ua.config_format == "yaml":
+        yaml_data = read_config(ua.yaml_file)
+        work_dir = yaml_data["work dir"]
+
+    elif ua.config_format == "cli":
+        work_dir = ua.work_dir
+
+    # Collect all the data
+    logger.info("Starting training process.")
+    if ua.hft:
+        type_list = ["aft", "hft"]
+    else:
+        type_list = ["aft"]
+    for data_type in type_list:
+        ids, ts_data, sweep_types, sel_coeffs = get_data(ua.training_data, data_type)
+        lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+        # Convert to numerical ohe IDs
+        num_ids = np.array([lab_dict[lab] for lab in ids])
+        ohe_ids = to_categorical(
+            np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+        )
+
+        class_weights = dict(
+            enumerate(
+                compute_class_weight(
+                    class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+                )
+            )
+        )
+        print(f"Class weights: {class_weights}")
+
+        if data_type == "aft":
+            # Needs to be in correct dims order for Conv1D layer
+            datadim = ts_data.shape[1:]
+            logger.info(
+                f"TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+            )
+        else:
+            logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+            datadim = ts_data.shape[1:]
+            logger.info(f"{len(ts_data)} samples in dataset.")
+
+        logger.info("Splitting Partition")
+        (
+            ts_train_data,
+            ts_val_data,
+            ts_test_data,
+            train_labs,
+            val_labs,
+            test_labs,
+            train_s,
+            val_s,
+            test_s,
+        ) = split_partitions(ts_data, ohe_ids, sel_coeffs)
+
+        # Time-series model training and evaluation
+        logger.info("Training time-series model.")
+
+
+        # scale data to the range of [0, 1]
+        trainData = trainData.astype("float32") / 255.0
+        testData = testData.astype("float32") / 255.0
+
+        # wrap our model into a scikit-learn compatible classifier
+        print("[INFO] initializing model...")
+
+        # define a grid of the hyperparameter search space
+        kernel_sizeLayerer_a = [16, 32, 64, 128, 256]
+        kernel_sizeLayerer_b = [16, 32, 64, 128, 256]
+        dropoutCovLayer_a= [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        # learnRate = [1e-2, 1e-3, 1e-4]
+        add_one_layer= [False, True]
+        # add_two_layer= [False, True]
+        dropoutFullLayer_a = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        dropoutFullLayer_b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        dropoutFullLayer_c = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
+        # batchSize = [4, 8, 16, 32]
+        # epochs = [10, 20, 30, 40]
+
+
+        # create a dictionary from the hyperparameter grid
+        grid = dict(
+        	kernel_sizeLayerer_a=kernel_sizeLayerer_a,
+        	kernel_sizeLayerer_b=kernel_sizeLayerer_b,
+        	dropoutCovLayer_a=dropoutCovLayer_a,
+        	add_one_layer=add_one_layer,
+        	dropoutFullLayer_a=dropoutFullLayer_a,
+            dropoutFullLayer_b=dropoutFullLayer_b,
+        	dropoutFullLayer_c=dropoutFullLayer_c
+        )
+
+        param_grid = ParameterGrid(grid)
+
+        for params in param_grid:
+            cur_model = get_mlp_model(**params)
+
+            # print(model.summary())
+            plot_model(
+                cur_model, to_file="model_plot.png", show_shapes=True, show_layer_names=True
+            )
+            cur_trained_model = fit_model(
+                work_dir,
+                cur_model,
+                data_type,
+                ts_train_data,
+                train_labs,
+                train_s,
+                ts_val_data,
+                val_labs,
+                val_s,
+                class_weights,
+                ua.experiment_name,
+            )
+            evaluate_model( ### need to modify to store fitted results
+                cur_trained_model,
+                ts_test_data,
+                test_labs,
+                test_s,
+                work_dir,
+                ua.experiment_name,
+                data_type,
+                lab_dict,
+            )
+
+
+        logger.info(f"SP Data shape (samples, haps): {sp_train_data.shape}")
+        # Use only the final timepoint
+        sp_train_data = np.squeeze(ts_train_data[:, -1, :])
+        sp_val_data = np.squeeze(ts_val_data[:, -1, :])
+        sp_test_data = np.squeeze(ts_test_data[:, -1, :])
+        logger.info(f"SP Data shape (samples, haps): {sp_train_data.shape}")
+
+        sp_datadim = sp_train_data.shape[-1]
+
+        for params in param_grid:
+            cur_model = model.__init__(self, **params)
+
+            sp_datadim = sp_train_data.shape[-1]
+            model = create_1Samp_model(sp_datadim, len(lab_dict))
+
+            trained_model = fit_model(
+                work_dir,
+                model,
+                data_type,
+                sp_train_data,
+                train_labs,
+                train_s,
+                sp_val_data,
+                val_labs,
+                val_s,
+                class_weights,
+                ua.experiment_name,
+            )
+            evaluate_model(
+                trained_model,
+                sp_test_data,
+                test_labs,
+                test_s,
+                work_dir,
+                ua.experiment_name,
+                data_type,
+                lab_dict,
+            )
+            )
+
+
+
+# # initialize a random search with a 3-fold cross-validation and then
+# # start the hyperparameter search process
+# print("[INFO] performing grid search...")
+# searcher = GridSearchCV(estimator=model, n_jobs=-1, cv=1
+# 	param_distributions=grid, scoring="accuracy")
+# searchResults = searcher.fit(train_data, [train_labs, train_s])
+#
+#
+# # summarize grid search information
+# bestScore = searchResults.best_score_
+# bestParams = searchResults.best_params_
+# print("[INFO] best score is {:.2f} using {}".format(bestScore,
+# 	bestParams))
+#
+# # extract the best model, make predictions on our data, and show a
+# # classification report
+# print("[INFO] evaluating the best model...")
+# bestModel = searchResults.best_estimator_
+# accuracy = bestModel.score(testData, testLabels)
+# print("accuracy: {:.2f}%".format(accuracy * 100))
diff --git a/timesweeper/parse_slim_logs.py b/timesweeper/parse_slim_logs.py
new file mode 100644
index 0000000..c2db182
--- /dev/null
+++ b/timesweeper/parse_slim_logs.py
@@ -0,0 +1,109 @@
+from glob import glob
+import os
+
+import pandas as pd
+from tqdm import tqdm
+
+from timesweeper.utils.gen_utils import read_config
+
+
+def parse_cmd(cmd):
+    """Create dictionary with all information in the bash command running slim."""
+    cmd = [i for i in cmd.split() if i not in ["slim", "-d"]]
+
+    cmd_dict = {}
+    for i in cmd:
+        if "=" in i:
+            clean_i = i.split("=")
+            cmd_dict[clean_i[0].replace("'", "").replace('"', "")] = (
+                clean_i[1].replace("'", "").replace('"', "")
+            )
+
+    return cmd_dict
+
+
+def get_samp_gens(loglist):
+    """Get sampling generations."""
+    samp_entries = [i for i in loglist if "Sampling at generation" in i]
+    samp_entries = set(samp_entries)
+    samp_gens = [i.split()[-1] for i in samp_entries]
+
+    return sorted([int(i) for i in samp_gens])
+
+
+def count_restarts(loglist):
+    restarts = 0
+    for i in loglist:
+        if "RESTARTING" in i:
+            restarts += 1
+
+    return restarts
+
+
+def track_sel_freq(loglist, num_samps):
+    freq_entries = [i for i in loglist if "SEGREGATING" in i and "FIXED" not in i]
+    if len(freq_entries) > 0:
+        freqs = [float(i.split()[-1]) for i in freq_entries]
+        return freqs[len(freqs) - num_samps :]
+
+    else:
+        return [0.0] * num_samps
+
+
+def get_rep_from_filename(filename):
+    rep = filename.split("/")[-1].split(".")[0]
+    return rep
+
+
+def parse_logfile(logfile):
+    with open(logfile, "r") as ifile:
+        loglist = [i.strip() for i in ifile.readlines()]
+
+    log_dict = parse_cmd(loglist[0])
+    log_dict["sampGens"] = get_samp_gens(loglist)
+    log_dict["selAlleleFreq"] = track_sel_freq(loglist, len(log_dict["sampGens"]))
+    log_dict["numRestarts"] = count_restarts(loglist)
+    log_dict["rep"] = get_rep_from_filename(log_dict["outFileVCF"])
+    log_dict["sampOffset"] = 10000 - log_dict["sampGens"][0]
+    for i in ["outFileVCF", "outFileMS", "dumpFile"]:
+        del log_dict[i]
+
+    return log_dict
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir, schema, threads = (
+        yaml_data["work dir"],
+        ua.experiment_name,
+        ua.threads,
+    )
+
+    logfiles = glob(f"{work_dir}/logs/*/*.log", recursive=True)
+
+    log_dict_list = []
+    for l in tqdm(logfiles, desc="Parsing logs"):
+        try:
+            log_dict_list.append(parse_logfile(l))
+        except Exception as e:
+            print(l, "couldn't work")
+
+    df = pd.DataFrame(log_dict_list)
+    df = df[
+        [
+            "rep",
+            "sweep",
+            "selCoeff",
+            "sampOffset",
+            "numRestarts",
+            "numSamples",
+            "seed",
+            "physLen",
+            "sampGens",
+            "selAlleleFreq",
+        ]
+    ]
+    df.to_csv(f"{work_dir}/{schema}_params.tsv", index=False, header=True, sep="\t")
+
+    for i in logfiles:
+        os.remove(i)
diff --git a/timesweeper/plotting/create_freq_track.py b/timesweeper/plotting/create_freq_track.py
index 483c4f8..3d5a0bd 100644
--- a/timesweeper/plotting/create_freq_track.py
+++ b/timesweeper/plotting/create_freq_track.py
@@ -1,5 +1,5 @@
-from ..utils import snp_utils as su
-from ..utils.gen_utils import get_logger, read_config
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import get_logger, read_config
 import numpy as np
 import allel
 import os
diff --git a/timesweeper/plotting/plot_bedfile_inputs.py b/timesweeper/plotting/plot_bedfile_inputs.py
index eec5f93..ae1d206 100644
--- a/timesweeper/plotting/plot_bedfile_inputs.py
+++ b/timesweeper/plotting/plot_bedfile_inputs.py
@@ -9,9 +9,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from find_sweeps_vcf import prep_ts_aft
-from .utils import snp_utils as su
-from .utils.gen_utils import read_config
+from timesweeper.utils import snp_utils as su
+from timesweeper.utils.gen_utils import read_config
 
 warnings.filterwarnings("ignore")
 
diff --git a/timesweeper/plotting/plot_training_data.py b/timesweeper/plotting/plot_training_data.py
index 58ffc84..5019e4d 100644
--- a/timesweeper/plotting/plot_training_data.py
+++ b/timesweeper/plotting/plot_training_data.py
@@ -36,11 +36,13 @@ def makeHeatmap(mat_type, data, plotTitle, axTitles, plotFileName):
             axes[i].set_title(axTitles[i], fontsize=14)
             axes[i].set_yticks([], minor=False)
             if "zoomed" in plotFileName:
-                axes[i].set_yticks([10.5], minor=True)
-                axes[i].set_yticklabels(["25"], minor=True)
+                # axes[i].set_yticks([1], minor=True)  # Single allele version
+                # axes[i].set_yticklabels(["1"], minor=True)  # Single allele version
+                axes[i].set_yticks([int(len(data[i]) / 2) + 0.51], minor=True)
+                axes[i].set_yticklabels(["Center"], minor=True)
             else:
-                axes[i].set_yticks([25.5], minor=True)
-                axes[i].set_yticklabels(["25"], minor=True)
+                axes[i].set_yticks([int(len(data[i]) / 2) + 0.5], minor=True)
+                axes[i].set_yticklabels(["Center"], minor=True)
 
             axes[i].set_xlabel("Timepoint")
             axes[i].set_xticks(list(range(data[i].shape[-1] + 1)))
@@ -56,6 +58,9 @@ def makeHeatmap(mat_type, data, plotTitle, axTitles, plotFileName):
         fig, axes = plt.subplots(1, len(data))
         for i in range(len(data)):
             heatmap = (axes[i].pcolor(data[i], cmap=plt.cm.Blues, norm=normscheme,),)[0]
+            # axes[i].set_yticks([1], minor=True)  # Single allele version
+            # axes[i].set_yticklabels(["1"], minor=True)  # Single allele version
+            # axes[i].set_yticks([int(len(data[i]) / 2) + 0.5], minor=True)
 
             plt.colorbar(heatmap, ax=axes[i])
             axes[i].set_title(axTitles[i], fontsize=14)
@@ -85,9 +90,14 @@ def readData(picklefile, data_type):
     sweep_types = pikl_dict.keys()
     data_dict = {}
     for key in sweep_types:
-        data_dict[key] = [
-            pikl_dict[key][rep][data_type.lower()] for rep in pikl_dict[key].keys()
-        ]
+        try:
+            data_dict[key] = [
+                pikl_dict[key][rep][data_type.lower()]
+                for rep in pikl_dict[key].keys()
+                if data_type.lower() in list(pikl_dict[key][rep].keys())
+            ]
+        except:
+            continue
 
     return data_dict
 
@@ -114,6 +124,7 @@ def main(ua):
     mat_types = get_mat_types(ua.input_pickle)
 
     for mat_type in mat_types:
+        print(mat_type)
         if mat_type == "sel_coeff":
             continue
 
@@ -152,12 +163,16 @@ def main(ua):
             )
             mean_diffs[lab] = mean_change
 
-            plt.plot(mean_change, label=lab.capitalize())
+            plt.plot(mean_change, label=lab.upper())
 
+        third_size = int(mean_data["neut"].shape[0] / 3)
         if mat_type == "aft":
             makeHeatmap(
                 mat_type,
-                [mean_data[i][15:36, :] for i in mean_data],
+                [mean_data[i][third_size : (2 * third_size), :] for i in mean_data],
+                # [
+                #    mean_data[i] for i in mean_data
+                # ],  # Use for single allele window size sims
                 schema_name,
                 [i.upper() for i in labs],
                 base_filename + f".zoomed.pdf",
@@ -174,7 +189,7 @@ def main(ua):
         elif mat_type == "hft":
             makeHeatmap(
                 mat_type,
-                [mean_data[i][:40, :] for i in mean_data],
+                [mean_data[i][:40] for i in mean_data],
                 schema_name,
                 [i.upper() for i in labs],
                 base_filename + f".zoomed.pdf",
@@ -182,7 +197,7 @@ def main(ua):
             for j in range(2):
                 makeHeatmap(
                     mat_type,
-                    [raw_data[i][j][:40, :] for i in raw_data],
+                    [raw_data[i][j] for i in raw_data],
                     schema_name,
                     [i.upper() for i in labs],
                     base_filename + f".{j}.single.pdf",
diff --git a/timesweeper/plotting/plot_windows.py b/timesweeper/plotting/plot_windows.py
index 44922ec..1cbe5b1 100644
--- a/timesweeper/plotting/plot_windows.py
+++ b/timesweeper/plotting/plot_windows.py
@@ -448,48 +448,8 @@ def main():
     # plot_means(datadict)
     plot_proportions(datadict)
     # plot_maxes(datadict)
-
     # plot_violinplots(datadict)
 
-    """
-    #These are jank, don't use before fixing
-    plot_roc(np.array(aft_trues), np.array(aft_probs), "aft_ROC", "aft_ROC.pdf", True)
-    plot_prec_recall(
-        np.array(aft_trues), np.array(aft_probs), "aft_PR", "aft_PR.pdf", True
-    )
-
-    aft_cm = confusion_matrix(aft_trues, aft_preds)
-    print(aft_cm)
-    plot_confusion_matrix(
-        ".",
-        aft_cm,
-        target_names=["Neutral", "SSV", "SDN"],  # , "Within 10kb"],
-        title="aft_Confmat_normed_foo",
-        normalize=True,
-    )
-    plot_confusion_matrix(
-        ".",
-        aft_cm,
-        target_names=["Neutral", "SSV", "SDN"],  # , "Within 10kb"],
-        title="aft_Confmat_unnormed_foo",
-        normalize=False,
-    )
-
-    plot_roc(np.array(hft_trues), np.array(hft_probs), "hft_ROC", "hft_ROC.pdf", True)
-    plot_prec_recall(
-        np.array(hft_trues), np.array(hft_probs), "hft_PR", "hft_PR.pdf", True
-    )
-    hft_cm = confusion_matrix(hft_trues, hft_preds)
-    print(hft_cm)
-    plot_confusion_matrix(
-        ".",
-        hft_cm,
-        target_names=["Neut", "SSV", "SDN"],
-        title="hft_Confmat",
-        normalize=True,
-    )
-    """
-
 
 if __name__ == "__main__":
     main()
diff --git a/timesweeper/plotting/plotting_utils.py b/timesweeper/plotting/plotting_utils.py
index 3b13ba8..381345f 100644
--- a/timesweeper/plotting/plotting_utils.py
+++ b/timesweeper/plotting/plotting_utils.py
@@ -179,14 +179,15 @@ def plot_roc(y_true, y_probs, schema, scenarios, outfile, combos=True, aggregate
     One v All ROC Curves for all scenarios.
     """
     labs = label_binarize(y_true, classes=list(range(len(scenarios))))
+    if labs.shape[1] == 1:
+        labs = np.hstack((1 - labs, labs))
+
     for i in range(len(scenarios)):
         fpr, tpr, thresh = roc_curve(labs[:, i], y_probs[:, i])
         swp_auc_val = auc(fpr, tpr)
-        plt.plot(
-            fpr, tpr, label=f"{scenarios[i].capitalize()} vs All: {swp_auc_val:.4}"
-        )
+        plt.plot(fpr, tpr, label=f"{scenarios[i].upper()} vs All: {swp_auc_val:.4}")
 
-    plt.title(f"ROC Curves {schema.capitalize()}")
+    plt.title(f"ROC Curves {schema.upper()}")
     plt.xlabel("FPR")
     plt.ylabel("TPR")
     plt.legend(loc="lower right")
@@ -202,17 +203,21 @@ def plot_prec_recall(
     """
     Identical to plot_roc_curves except it's precision/recall.
     """
+    # TODO Double check this
     labs = label_binarize(y_true, classes=list(range(len(scenarios))))
+    if labs.shape[1] == 1:
+        labs = np.hstack((1 - labs, labs))
+
     for i in range(len(scenarios)):
         prec, recall, thresh = precision_recall_curve(labs[:, i], y_probs[:, i])
         ap_val = average_precision_score(labs[:, i], y_probs[:, i])
         plt.plot(
             prec,
             recall,
-            label=f"{scenarios[i].capitalize()} vs All - Avg Precision: {ap_val:.4}",
+            label=f"{scenarios[i].upper()} vs All - Avg Precision: {ap_val:.4}",
         )
 
-    plt.title(f"PR Curve {schema.capitalize()}")
+    plt.title(f"PR Curve {schema.upper()}")
     plt.legend(loc="lower right")
     plt.xlabel("Recall")
     plt.ylabel("Precision")
@@ -225,9 +230,9 @@ def plot_sel_coeff_preds(true_class, s_true, s_pred, outfile, scenarios):
     """Plot s predictions against true value, color by scenario."""
     for g in np.unique(true_class):
         i = np.where(true_class == g)
-        plt.scatter(s_true[i], s_pred[i], label=scenarios[g].capitalize())
+        plt.scatter(s_true[i], s_pred[i], label=scenarios[g].upper())
         plt.annotate(
-            f"r^2 of {scenarios[g].capitalize()}: {np.round(r2_score(s_true[i], s_pred[i]), 2)}",
+            f"r^2 of {scenarios[g].upper()}: {np.round(r2_score(s_true[i], s_pred[i]), 2)}",
             (0.05, 0.27),
         )
 
diff --git a/timesweeper/process_vcfs.py b/timesweeper/process_vcfs.py
index 53495c7..2625911 100644
--- a/timesweeper/process_vcfs.py
+++ b/timesweeper/process_vcfs.py
@@ -1,19 +1,11 @@
-import logging
 import multiprocessing as mp
 import os
 import subprocess
 from glob import glob
 from itertools import cycle
-import random as rd
-from tqdm import tqdm
-
-from .utils.gen_utils import read_config
-
-logging.basicConfig()
-logger = logging.getLogger("vcf_processing")
-logger.setLevel("INFO")
 
 
+# VCF Processing
 def read_multivcf(input_vcf):
     """Reads in file and returns as list of strings."""
     with open(input_vcf, "r") as input_file:
@@ -35,20 +27,6 @@ def split_multivcf(vcf_lines, header):
     return split_vcfs
 
 
-def make_vcf_dir(input_vcf):
-    """Creates directory named after vcf basename."""
-    dirname = os.path.basename(input_vcf).split(".")[0]
-    dirpath = os.path.dirname(input_vcf)
-    vcf_dir = os.path.join(dirpath, dirname)
-    if os.path.exists(vcf_dir):
-        for ifile in glob(f"{vcf_dir}/*"):
-            os.remove(ifile)
-
-    os.makedirs(vcf_dir, exist_ok=True)
-
-    return vcf_dir
-
-
 def write_vcfs(vcf_lines, vcf_dir):
     """Writes list of vcf entries to numerically-sorted vcf files."""
     for idx, lines in enumerate(vcf_lines):
@@ -81,25 +59,56 @@ def index_vcf(vcf):
 
 def merge_vcfs(vcf_dir):
     num_files = len(glob(f"{vcf_dir}/*.vcf.sorted.gz"))
-    cmd = f"""bcftools merge -Ov \
-            --force-samples -0 \
-            {" ".join([f"{vcf_dir}/{i}.vcf.sorted.gz" for i in range(num_files)])} > \
-            {vcf_dir}/merged.vcf \
-            """
+    if num_files == 1:
+        cmd = f"""zcat {f"{vcf_dir}/0.vcf.sorted.gz"} > {vcf_dir}/merged.vcf"""
+
+    else:
+        cmd = f"""bcftools merge -Ov -0 \
+                --force-samples --info-rules 'MT:join,S:join' \
+                {" ".join([f"{vcf_dir}/{i}.vcf.sorted.gz" for i in range(num_files)])} > \
+                {vcf_dir}/merged.vcf \
+                """
+
     subprocess.run(cmd, shell=True)
 
 
+def get_num_inds(vcf_file):
+    num_ind = subprocess.check_output(
+        """awk '{if ($1 == "#CHROM"){print NF-9; exit}}' """ + vcf_file, shell=True,
+    )
+    return int(num_ind)
+
+
 def cleanup_intermed(vcf_dir):
     for ifile in glob(f"{vcf_dir}/*"):
-        if "merged" not in ifile:
+        if "merged" not in ifile and "final" not in ifile:
+            pass
             os.remove(ifile)
 
 
-def worker(input_vcf, num_tps, vcf_header, verbose=False):
+def make_vcf_dir(input_vcf):
+    """Creates directory named after vcf basename."""
+    dirname = os.path.basename(input_vcf).split(".")[0]
+    dirpath = os.path.dirname(input_vcf)
+    vcf_dir = os.path.join(dirpath, dirname)
+    if os.path.exists(vcf_dir):
+        for ifile in glob(f"{vcf_dir}/*"):
+            os.remove(ifile)
+
+    os.makedirs(vcf_dir, exist_ok=True)
+
+    final_vcf = f"{input_vcf}.final"
+    if os.path.exists(final_vcf):
+        os.rename(final_vcf, f"{vcf_dir}/{final_vcf.split('/')[-1]}")
+
+    return vcf_dir
+
+
+def process_vcfs(input_vcf, num_tps):
     try:
         # Split into multiples after SLiM just concats to same file
         raw_lines = read_multivcf(input_vcf)
-        split_lines = split_multivcf(raw_lines, vcf_header)
+        split_lines = split_multivcf(raw_lines, "##fileformat=VCFv4.2")
         if len(split_lines) > 0:
             split_lines = split_lines[len(split_lines) - num_tps :]
 
@@ -110,39 +119,18 @@ def worker(input_vcf, num_tps, vcf_header, verbose=False):
             # Now index and merge
             [index_vcf(vcf) for vcf in glob(f"{vcf_dir}/*.vcf")]
             merge_vcfs(vcf_dir)
+
             cleanup_intermed(vcf_dir)
-        else:
-            pass
 
     except Exception as e:
-        if verbose:
-            logger.warning(e)
+        print(f"[ERROR] Couldn't process {e}")
         pass
 
 
 def main(ua):
-    """Splits multivcf generated by SLiM using simulate.py and creates a merged/indexed vcf to create training data with."""
-    yaml_data = read_config(ua.yaml_file)
-    work_dir, samp_sizes, threads, vcf_header = (
-        yaml_data["work dir"],
-        yaml_data["sample sizes"],
-        ua.threads,
-        ua.vcf_header,
-    )
+    vcflist = glob(f"{ua.in_dir}/*/*/*.multivcf")
 
-    logger.info(f"Processing multiVCFs in {work_dir} using {threads} threads.")
-    # Quick method to sample from each class evenly
-    input_vcfs = []
-    for i in ["neut", "ssv", "sdn"]:
-        input_vcfs.append(rd.sample(glob(f"{work_dir}/vcfs/{i}/*.multivcf"), 10000))
-
-    pool = mp.Pool(threads)
-    pool.starmap(
-        worker,
-        tqdm(
-            zip(input_vcfs, cycle([len(samp_sizes)]), cycle([vcf_header])),
-            desc="Processing VCFs",
-            total=len(input_vcfs),
-        ),
-        chunksize=5,
-    )
+    with mp.Pool(ua.threads) as p:
+        p.map(vcflist, cycle([ua.num_tps]))
+
+    print("[INFO] Done")
diff --git a/timesweeper/simulate_custom.py b/timesweeper/simulate_custom.py
index b7d03e5..bfd2deb 100644
--- a/timesweeper/simulate_custom.py
+++ b/timesweeper/simulate_custom.py
@@ -1,16 +1,17 @@
-import argparse
 import logging
 import multiprocessing as mp
 import os
-import shutil
 import subprocess
 import sys
+from glob import glob
 
 import numpy as np
 import yaml
 
+from timesweeper import process_vcfs as pv
+
 logging.basicConfig()
-logger = logging.getLogger("simple_simulate")
+logger = logging.getLogger("sim_custom")
 
 
 def read_config(yaml_file):
@@ -21,7 +22,8 @@ def read_config(yaml_file):
     return yamldata
 
 
-def randomize_selCoeff_uni(lower_bound=0.0, upper_bound=0.25):
+# Simulation
+def randomize_selCoeff_uni(lower_bound=0.00025, upper_bound=0.25):
     """Draws selection coefficient from log uniform dist to vary selection strength."""
     rng = np.random.default_rng(
         np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
@@ -34,9 +36,13 @@ def randomize_sampGens(num_timepoints, dev=50, span=200):
     rng = np.random.default_rng(
         np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
     )
-
     start = round(rng.uniform(-dev, dev, 1)[0])
-    sampGens = [round(i) for i in np.linspace(start, start + span + 1, num_timepoints)]
+    if num_timepoints == 1:
+        sampGens = [start + span]
+    else:
+        sampGens = [
+            round(i) for i in np.linspace(start, start + span + 1, num_timepoints)
+        ]
 
     return sampGens
 
@@ -58,7 +64,6 @@ def make_d_block(
     Please note that when feeding strings as a constant you must escape them since this is a shell process.
     """
     selCoeff = randomize_selCoeff_uni()
-
     sampGens = [str(i) for i in randomize_sampGens(num_sample_points)]
 
     d_block = f"""\
@@ -79,8 +84,10 @@ def make_d_block(
     return d_block
 
 
-def run_slim(slimfile, slim_path, d_block, logfile):
-    cmd = " ".join([slim_path, d_block, slimfile, ">>", logfile]).replace("    ", "")
+def simulate(slim_path, d_block, slimfile, logfile):
+    cmd = " ".join(["time", slim_path, d_block, slimfile, ">>", logfile]).replace(
+        "    ", ""
+    )
     with open(logfile, "w") as ofile:
         ofile.write(cmd)
 
@@ -89,8 +96,15 @@ def run_slim(slimfile, slim_path, d_block, logfile):
     except subprocess.CalledProcessError as e:
         logger.error(e.output)
 
-    sys.stdout.flush()
-    sys.stderr.flush()
+
+def simulate_prep(
+    vcf_file, num_sample_points, slimfile, slim_path, d_block, logfile, dumpFile
+):
+    simulate(slim_path, d_block, slimfile, logfile)
+    os.remove(dumpFile)
+
+    pv.process_vcfs(vcf_file, num_sample_points)
+    os.remove(vcf_file)
 
 
 def main(ua):
@@ -99,18 +113,18 @@ def main(ua):
     Currently only works with 1 pop models where m2 is the sweep mutation.
     Otherwise everything else is identical to stdpopsim version, just less complex.
 
-    Generalized block of '-d' arguments to give to SLiM at the command line allow for 
+    Generalized block of '-d' arguments to give to SLiM at the command line allow for
     flexible script writing within the context of this wrapper. If you write your SLiM script
-    to require args set at runtime, this should be easily modifiable to do what you need and 
+    to require args set at runtime, this should be easily modifiable to do what you need and
     get consistent results to plug into the rest of the workflow.
 
     The two things you *will* need to specify in your '-d' args to SLiM (and somewhere in the slim script) are:
-    - sweep [str] One of "neut", "sdn", or "ssv". If you're testing only a neut/sdn model, 
+    - sweep [str] One of "neut", "sdn", or "ssv". If you're testing only a neut/sdn model,
         make the ssv a dummy switch for the neutral scenario.
-    - outFile [path] You will need to define this as a population outputVCFSample input, with replace=F and append=T. 
+    - outFile [path] You will need to define this as a population outputVCFSample input, with replace=F and append=T.
         This does *not* need to be specified by you in the custom -d block, it will be standardized to work with the rest of the pipeline using work_dir.
         example line for slim script: `p1.outputVCFSample(sampleSizePerStep, replace=F, append=T, filePath=outFile);`
-        
+
     """
     yaml_data = read_config(ua.yaml_file)
     (
@@ -140,7 +154,7 @@ def main(ua):
 
     sweeps = ["neut", "sdn", "ssv"]
 
-    for i in [vcf_dir, ms_dir, dumpfile_dir, logfile_dir]:
+    for i in [vcf_dir, dumpfile_dir, logfile_dir]:
         for sweep in sweeps:
             os.makedirs(f"{i}/{sweep}", exist_ok=True)
 
@@ -158,31 +172,28 @@ def main(ua):
             dumpFile = f"{dumpfile_dir}/{sweep}/{rep}.dump"
             logFile = f"{logfile_dir}/{sweep}/{rep}.log"
 
-            # Grab those constants to feed to SLiM
-            if rep == 0:
-                d_block = make_d_block(
-                    sweep,
+            d_block = make_d_block(
+                sweep,
+                outFileVCF,
+                outFileMS,
+                dumpFile,
+                num_sample_points,
+                inds_per_tp,
+                physLen,
+                False,
+            )
+
+            mp_args.append(
+                (
                     outFileVCF,
-                    outFileMS,
-                    dumpFile,
                     num_sample_points,
-                    inds_per_tp,
-                    physLen,
-                    True,
-                )
-            else:
-                d_block = make_d_block(
-                    sweep,
-                    outFileVCF,
-                    outFileMS,
+                    slim_file,
+                    slim_path,
+                    d_block,
+                    logFile,
                     dumpFile,
-                    num_sample_points,
-                    inds_per_tp,
-                    physLen,
-                    False,
                 )
-
-            mp_args.append((slim_file, slim_path, d_block, logFile))
+            )
 
     pool = mp.Pool(processes=ua.threads)
-    pool.starmap(run_slim, mp_args, chunksize=1)
+    pool.starmap(simulate_prep, mp_args, chunksize=1)
diff --git a/timesweeper/simulate_stdpopsim.py b/timesweeper/simulate_stdpopsim.py
index 6e6215d..8ad9519 100644
--- a/timesweeper/simulate_stdpopsim.py
+++ b/timesweeper/simulate_stdpopsim.py
@@ -1,15 +1,18 @@
 import argparse as ap
 import logging
 import multiprocessing as mp
-import os, sys
+import os
 import re
 import subprocess
+import sys
+from glob import glob
 from itertools import cycle
 
 import numpy as np
 import pandas as pd
 
-from .utils.gen_utils import read_config
+from timesweeper import simulate_custom as sc
+from timesweeper.utils.gen_utils import read_config
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
@@ -156,7 +159,7 @@ def make_sel_blocks(sweep, sel_gen, pop, dumpFileName):
     \n{sel_gen} late(){{
         if (sweep == "ssv")
         {{
-            muts = sim.mutationssvype(m1);
+            muts = sim.mutationsOfType(m1);
             if (size(muts))
             {{
                 mut = NULL;
@@ -205,7 +208,7 @@ def make_sel_blocks(sweep, sel_gen, pop, dumpFileName):
             }}
             else
             {{
-                muts = sim.mutationssvype(m2);
+                muts = sim.mutationsOfType(m2);
                 if (size(muts) == 0)
                 {{
                     print("LOST at gen " + sim.generation + " - RESTARTING");
@@ -249,7 +252,7 @@ def randomize_selCoeff_loguni(lower_bound=0.005, upper_bound=0.5):
     return 10 ** rand_log[0]
 
 
-def randomize_selCoeff_uni(lower_bound=0.0, upper_bound=0.05):
+def randomize_selCoeff_uni(lower_bound=0.00025, upper_bound=0.25):
     """Draws selection coefficient from log uniform dist to vary selection strength."""
     rng = np.random.default_rng(
         np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
@@ -298,6 +301,29 @@ def run_slim(slimfile, slim_path):
     sys.stderr.flush()
 
 
+def process_vcfs(input_vcf, num_tps):
+    try:
+        # Split into multiples after SLiM just concats to same file
+        raw_lines = sc.read_multivcf(input_vcf)
+        split_lines = sc.split_multivcf(raw_lines, "##fileformat=VCFv4.2")
+        if len(split_lines) > 0:
+            split_lines = split_lines[len(split_lines) - num_tps :]
+
+            # Creates subdir for each rep
+            vcf_dir = sc.make_vcf_dir(input_vcf)
+            sc.write_vcfs(split_lines, vcf_dir)
+
+            # Now index and merge
+            [sc.index_vcf(vcf) for vcf in glob(f"{vcf_dir}/*.vcf")]
+            sc.merge_vcfs(vcf_dir)
+
+            sc.cleanup_intermed(vcf_dir)
+
+    except Exception as e:
+        print(f"[ERROR] Couldn't process {e}")
+        pass
+
+
 def clean_args(ua):
     yaml_data = read_config(ua.yaml_file)
     (
@@ -408,7 +434,7 @@ def main(ua):
             Q, gen_time, max_years_b0, burn_in_gens, physLen = get_slim_info(raw_lines)
 
             # Pull from variable time of selection before sampling to make more robust
-            rand_sel_gen = randomize_selTime(sel_gen, 50 / Q)
+            rand_sel_gen = randomize_selTime(sel_gen, 200 / Q)
 
             burn_in_gens = int(round(burn_in_gens / Q))
 
@@ -516,9 +542,12 @@ def main(ua):
     for script in script_list:
         run_slim(script, slim_path)
 
-    # Cleanup
+    # Process VCFs and Cleanup
     for rep in replist:
         for sweep in sweeps:
+            vcf_file = f"{vcf_dir}/{sweep}/{rep}.multivcf"
+            process_vcfs(vcf_file, len(sample_sizes))
+            os.remove(vcf_file)
             dumpFile = f"{dumpfile_dir}/{sweep}/{rep}.dump"
             os.remove(dumpFile)
 
diff --git a/timesweeper/tp1_model.py b/timesweeper/tp1_model.py
new file mode 100644
index 0000000..9c40e21
--- /dev/null
+++ b/timesweeper/tp1_model.py
@@ -0,0 +1,708 @@
+"""
+This is a modified version of the `train_nets.py` module in the core Timesweeper package. 
+Modifications include:
+- Fully connected network instead of 1DCNN
+"""
+
+import logging
+import os
+import pickle
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    id_list = []
+    rep_list = []
+    data_list = []
+    sel_coeffs = []
+    sweep_types = []
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+    for sweep in pikl_dict.keys():
+        sweep_types.append(sweep)
+        for rep in pikl_dict[sweep].keys():
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+            id_list.append(sweep)
+            rep_list.append(rep)
+            sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+
+    return (
+        id_list,
+        rep_list,
+        np.stack(data_list),
+        sweep_types,
+        np.array(sel_coeffs).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir, model, data_type, train_data, train_s, val_data, val_s, experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    ts_data = np.squeeze(ts_data)
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    class_model = models.create_1tp_class_model(datadim, len(lab_dict))  # type: ignore
+    reg_model = models.create_1tp_reg_model(datadim)  # type: ignore
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = False
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/timesweeper/train_nets.py b/timesweeper/train_nets.py
new file mode 100644
index 0000000..0640aaf
--- /dev/null
+++ b/timesweeper/train_nets.py
@@ -0,0 +1,752 @@
+import logging
+import os
+import pickle
+import random
+import sys
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    id_list = []
+    rep_list = []
+    data_list = []
+    sel_coeffs = []
+    sweep_types = []
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+    for sweep in pikl_dict.keys():
+        sweep_types.append(sweep)
+        for rep in pikl_dict[sweep].keys():
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+
+            id_list.append(sweep)
+            rep_list.append(rep)
+            sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+
+    return (
+        id_list,
+        rep_list,
+        np.stack(data_list),
+        sweep_types,
+        np.array(sel_coeffs).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir, model, data_type, train_data, train_s, val_data, val_s, experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    """
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+    
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+    """
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_prob"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "2dcnn":
+        ts_train_data = np.expand_dims(ts_train_data, -1)
+        ts_val_data = np.expand_dims(ts_val_data, -1)
+        ts_test_data = np.expand_dims(ts_test_data, -1)
+        datadim = ts_train_data.shape[1:]
+        class_model = models.create_2D_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_2D_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "chonk":
+        class_model = models.create_big_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_big_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "rnn":
+        class_model = models.create_rnn_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_rnn_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        ts_train_data = np.expand_dims(ts_train_data, -1)
+        ts_val_data = np.expand_dims(ts_val_data, -1)
+        ts_test_data = np.expand_dims(ts_test_data, -1)
+        datadim = ts_train_data.shape[:-1]
+
+        class_model = models.create_transformer_class_model(
+            input_shape=datadim,
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=datadim,
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                pickle.dump(
+                    mm_scaler,
+                    open(f"{work_dir}/{ua.experiment_name}_selcoeff_scaler.pkl", "wb"),
+                )
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = False
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/timesweeper/train_nets_shic.py b/timesweeper/train_nets_shic.py
new file mode 100644
index 0000000..bb6d6fb
--- /dev/null
+++ b/timesweeper/train_nets_shic.py
@@ -0,0 +1,713 @@
+import logging
+import os
+import pickle
+import random
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
+from sklearn.utils import compute_class_weight
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
+
+from timesweeper import models
+
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+logging.basicConfig()
+logger = logging.getLogger("nets")
+logger.setLevel("INFO")
+
+seed = np.random.randint(1, 1e6)
+random.seed(seed)
+np.random.seed(seed)
+tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
+
+
+def scale_sel_coeffs(sel_coef_arr):
+    """Scale training data, return scaled data and scaler to use for test data."""
+    scaler = MinMaxScaler().fit(sel_coef_arr)
+
+    return scaler
+
+
+def get_data(input_pickle, data_type):
+    """
+    Loads data from pickle file and returns as list of labels and data.
+
+    Args:
+        input_pickle (str): Path to pickle created with make_training_features module.
+        data_type (str): Determines either hfs or aft files to search for.
+
+    Returns:
+        list[str]: List of sweep labels for each sample
+        np.arr: Array with all data stacked.
+    """
+    pikl_dict = pickle.load(open(input_pickle, "rb"))
+
+    return (
+        pikl_dict["sweep"],
+        pikl_dict["rep"],
+        np.stack(pikl_dict["data"]),
+        ["neut", "sdn", "ssv"],
+        np.array(pikl_dict["selcoeff"]).reshape(-1, 1),
+    )
+
+
+def split_partitions(data, labs, sel_coeffs, reps):
+    """
+        Splits all data and labels into partitions for train/val/testing.
+
+    Args:
+        data (np.arr): Data for training model.
+        labs (List): List of numeric labels for IDs
+        sel_coeffs (List[float]): List of selection coefficients for each rep
+
+    Returns:
+        Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
+    """
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+    ) = train_test_split(data, labs, sel_coeffs, reps, stratify=labs, test_size=0.3)
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+    ) = train_test_split(
+        val_data, val_labs, val_s, val_reps, stratify=val_labs, test_size=0.5
+    )
+    return (
+        train_data,
+        val_data,
+        test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        train_reps,
+        val_reps,
+        test_reps,
+    )
+
+
+def fit_class_model(
+    out_dir,
+    model,
+    data_type,
+    train_data,
+    train_labs,
+    val_data,
+    val_labs,
+    experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_accuracy"
+    train_target = train_labs
+    val_target = val_labs
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.1,
+        patience=20,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_class_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def fit_reg_model(
+    out_dir, model, data_type, train_data, train_s, val_data, val_s, experiment_name,
+):
+    """
+    Fits a given model using training/validation data, plots history after done.
+
+    Args:
+        out_dir (str): Base directory where data is located, model will be saved here.
+        model (Model): Compiled Keras model.
+        data_type (str): Whether data is HFS or aft data.
+        train_data (np.arr): Training data.
+        train_labs (list[int]): OHE labels for training set.
+        val_data (np.arr): Validation data.
+        val_labs (list[int]): OHE labels for validation set.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+
+    Returns:
+        Model: Fitted Keras model, ready to be used for accuracy characterization.
+    """
+    monitor = "val_mse"
+    train_target = train_s
+    val_target = val_s
+
+    if not os.path.exists(os.path.join(out_dir, "images")):
+        os.makedirs(os.path.join(out_dir, "images"), exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_dir, "trained_models")):
+        os.makedirs(os.path.join(out_dir, "trained_models"), exist_ok=True)
+
+    checkpoint = ModelCheckpoint(
+        os.path.join(out_dir, "trained_models", f"{model.name}_{data_type}"),
+        monitor=monitor,
+        verbose=1,
+        save_best_only=True,
+        save_weights_only=True,
+        mode="auto",
+    )
+
+    earlystop = EarlyStopping(
+        monitor=monitor,
+        min_delta=0.01,
+        patience=40,
+        verbose=1,
+        mode="auto",
+        restore_best_weights=True,
+    )
+
+    callbacks_list = [earlystop, checkpoint]
+
+    history = model.fit(
+        x=train_data,
+        y=train_target,
+        epochs=100,
+        verbose=2,
+        callbacks=callbacks_list,
+        validation_data=(val_data, val_target),
+        # class_weight=class_weights,
+    )
+
+    pu.plot_reg_training(
+        os.path.join(out_dir, "images"),
+        history,
+        f"{experiment_name}_{model.name}_{data_type}",
+    )
+
+    # Won't checkpoint handle this?
+    save_model(
+        model,
+        os.path.join(
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
+        ),
+    )
+
+    return model
+
+
+def evaluate_reg_model(
+    model,
+    test_data,
+    test_labs,
+    trans_test_s,
+    test_reps,
+    s_scaler,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
+    trues = np.argmax(test_labs, axis=1)
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
+
+    pred_dict = {
+        "rep": test_reps,
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        trues,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+    # Linear regression to correct for prediction boundary
+    (
+        train_pred_s,
+        test_pred_s,
+        train_true_s,
+        test_true_s,
+        train_corr_rep,
+        test_corr_rep,
+        train_corr_class,
+        test_corr_class,
+    ) = train_test_split(trans_pred_s, trans_test_s, test_reps, trues, test_size=0.3)
+
+    linreg = LinearRegression()
+    linreg.fit(train_pred_s, train_true_s)
+    logger.info(
+        f"\nScore for s correction linear model: {linreg.score(train_pred_s, train_true_s)}"
+    )
+
+    corrected_preds = linreg.predict(test_pred_s)
+    logger.info(
+        f"\nMSE for  s correction linear model: {mean_squared_error(test_true_s, corrected_preds)}"
+    )
+
+    if "log" in experiment_name:
+        pred_s = 10 ** (-corrected_preds)
+        test_s = 10 ** (-test_true_s)
+    elif "minmax" in experiment_name:
+        pred_s = s_scaler.inverse_transform(corrected_preds.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(test_true_s.reshape(-1, 1))
+    else:
+        pred_s = corrected_preds
+        test_s = test_true_s
+
+    pred_dict = {
+        "rep": test_corr_rep,
+        "class": [str_lab_dict[i] for i in test_corr_class],
+        "true_sel_coeff": test_s.flatten(),
+        "corrected_pred_sel_coeff": pred_s.flatten(),
+    }
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeff_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    )
+    pu.plot_sel_coeff_preds(
+        test_corr_class,
+        test_s,
+        pred_s,
+        os.path.join(
+            out_dir,
+            "images",
+            f"{experiment_name}_{model.name}_{data_type}_corrected_selcoeffs.pdf",
+        ),
+        scenarios,
+    )
+
+
+def evaluate_class_model(
+    model,
+    test_data,
+    test_labs,
+    test_reps,
+    out_dir,
+    scenarios,
+    experiment_name,
+    data_type,
+    lab_dict,
+):
+    """
+    Evaluates model using confusion matrices and plots results.
+
+    Args:
+        model (Model): Fit Keras model.
+        test_data (List[narr]): Testing data.
+        test_labs (narr): Testing labels.
+        test_s (narr): Selection coefficients to test against.
+        out_dir (str): Base directory data is located in.
+        scenarios (list[str]): Scenarios defined in config.
+        experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
+        data_type (str): Whether data is aft or hfs.
+    """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+
+    class_probs = model.predict(test_data)
+    class_predictions = np.argmax(class_probs, axis=1)
+    trues = np.argmax(test_labs, axis=1)
+
+    # Cannot for the life of me figure out why memory is shared b/t functions and this
+    # So it gets casted twice to break that chain
+    roc_trues = np.array(list(trues))
+    pr_trues = np.array(list(trues))
+
+    pred_dict = {
+        "rep": test_reps,
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
+    }
+    for str_lab in lab_dict:
+        pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
+
+    pred_df = pd.DataFrame(pred_dict)
+
+    os.makedirs(os.path.join(out_dir, "test_predictions"), exist_ok=True)
+    pred_df.to_csv(
+        os.path.join(
+            out_dir,
+            "test_predictions",
+            f"{experiment_name}_{model.name}_{data_type}_class_test_predictions.csv",
+        ),
+        header=True,
+        index=False,
+    )
+
+    lablist = [i.upper() for i in lab_dict]
+
+    conf_mat = confusion_matrix(trues, class_predictions)
+
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_normed",
+        normalize=True,
+    )
+    pu.plot_confusion_matrix(
+        os.path.join(out_dir, "images"),
+        conf_mat,
+        lablist,
+        title=f"{experiment_name}_{model.name}_{data_type}_confmat_unnormed",
+        normalize=False,
+    )
+
+    pu.print_classification_report(trues, class_predictions)
+
+    pu.plot_roc(
+        roc_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_roc.pdf"
+        ),
+    )
+
+    pu.plot_prec_recall(
+        pr_trues,
+        class_probs,
+        f"{experiment_name}_{model.name}_{data_type}",
+        scenarios,
+        os.path.join(
+            out_dir, "images", f"{experiment_name}_{model.name}_{data_type}_pr.pdf"
+        ),
+    )
+
+
+def main(ua):
+    yaml_data = read_config(ua.yaml_file)
+    work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
+
+    # Collect all the data
+    logger.info("Starting TP-SHIC training process.")
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs = get_data(
+        ua.training_data, ua.data_type
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        ts_data, _, ohe_ids, _, sel_coeffs, _, reps, _ = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
+    else:
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
+        )
+    )
+    logger.info(f"Class weights: {class_weights}")
+
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+        )
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
+
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
+
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
+            ts_train_data,
+            train_labs,
+            ts_val_data,
+            val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
+            test_labs,
+            test_reps,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape}"
+            )
+            train_idxs = np.where(
+                (train_s.flatten() > 0.0) & (train_labs[:, idx] == 1)
+            )[0]
+            val_idxs = np.where((val_s.flatten() > 0.0) & (val_labs[:, idx] == 1))[0]
+            test_idxs = np.where((test_s.flatten() > 0.0) & (test_labs[:, idx] == 1))[0]
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s[train_idxs])
+                vvals = -np.log10(val_s[val_idxs])
+                tevals = -np.log10(test_s[test_idxs])
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s[train_idxs])
+                trvals = mm_scaler.transform(train_s[train_idxs])
+                vvals = mm_scaler.transform(val_s[val_idxs])
+                tevals = mm_scaler.transform(test_s[test_idxs])
+            else:
+                mm_scaler = None
+                trvals = train_s[train_idxs]
+                vvals = val_s[val_idxs]
+                tevals = test_s[test_idxs]
+
+            plot = True
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s[train_idxs],
+                    ts_train_data[train_idxs, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[train_idxs, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
+
+            trained_reg_model = fit_reg_model(
+                work_dir,
+                reg_model,
+                ua.data_type,
+                ts_train_data[train_idxs],
+                trvals,
+                ts_val_data[val_idxs],
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
+            )
+            evaluate_reg_model(
+                trained_reg_model,
+                ts_test_data[test_idxs],
+                test_labs[test_idxs],
+                tevals,
+                [test_reps[i] for i in list(test_idxs)],
+                mm_scaler,
+                work_dir,
+                yaml_data["scenarios"],
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
+                lab_dict,
+            )
diff --git a/timesweeper/nets.py b/timesweeper/train_shoulder_nets.py
similarity index 55%
rename from timesweeper/nets.py
rename to timesweeper/train_shoulder_nets.py
index f02a902..e5369fc 100644
--- a/timesweeper/nets.py
+++ b/timesweeper/train_shoulder_nets.py
@@ -8,18 +8,19 @@
 import numpy as np
 import pandas as pd
 import tensorflow as tf
-from keras.utils import to_categorical
-from sklearn.metrics import confusion_matrix, mean_absolute_error
+from tensorflow.keras.utils import to_categorical
+from sklearn.metrics import confusion_matrix, mean_absolute_error, mean_squared_error
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import MinMaxScaler
+from sklearn.linear_model import LinearRegression
 from sklearn.utils import compute_class_weight
-from keras.callbacks import EarlyStopping, ModelCheckpoint
-from keras.models import save_model
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
+from tensorflow.keras.models import save_model
 
-from . import models
+from timesweeper import models
 
-from .plotting import plotting_utils as pu
-from .utils.gen_utils import read_config
+from timesweeper.plotting import plotting_utils as pu
+from timesweeper.utils.gen_utils import read_config
 
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
@@ -27,10 +28,11 @@
 logger = logging.getLogger("nets")
 logger.setLevel("INFO")
 
-seed = 42
+seed = np.random.randint(1, 1e6)
 random.seed(seed)
 np.random.seed(seed)
 tf.random.set_seed(seed)
+logger.info(f"Random seed is: {seed}")
 
 
 def scale_sel_coeffs(sel_coef_arr):
@@ -40,7 +42,7 @@ def scale_sel_coeffs(sel_coef_arr):
     return scaler
 
 
-def get_data(input_pickle, data_type):
+def get_data(input_pickle, data_type, win_size):
     """
     Loads data from pickle file and returns as list of labels and data.
 
@@ -53,26 +55,44 @@ def get_data(input_pickle, data_type):
         np.arr: Array with all data stacked.
     """
     id_list = []
+    rep_list = []
     data_list = []
     sel_coeffs = []
     sweep_types = []
+    offsets = []
     pikl_dict = pickle.load(open(input_pickle, "rb"))
     for sweep in pikl_dict.keys():
         sweep_types.append(sweep)
         for rep in pikl_dict[sweep].keys():
-            id_list.append(sweep)
-            data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            try:
+                data_list.append(np.array(pikl_dict[sweep][rep][data_type.lower()]))
+            except:
+                continue
+
+            if sweep != "neut":
+                if abs(pikl_dict[sweep][rep]["center_offset"]) > 0:
+                    id_list.append("neut")
+                else:
+                    id_list.append(sweep)
+
+            else:
+                id_list.append(sweep)
+
+            rep_list.append(rep)
             sel_coeffs.append(pikl_dict[sweep][rep]["sel_coeff"])
+            offsets.append(pikl_dict[sweep][rep]["center_offset"])
 
     return (
         id_list,
+        rep_list,
         np.stack(data_list),
         sweep_types,
         np.array(sel_coeffs).reshape(-1, 1),
+        np.array(offsets).reshape(-1, 1),
     )
 
 
-def split_partitions(data, labs, sel_coeffs):
+def split_partitions(data, labs, sel_coeffs, reps, offsets):
     """
         Splits all data and labels into partitions for train/val/testing.
 
@@ -84,11 +104,39 @@ def split_partitions(data, labs, sel_coeffs):
     Returns:
         Tuple[List[narr], List[narr], List[narr], List[int], List[int], List[int]]: Train/val/test splits of IDs and labs
     """
-    train_data, val_data, train_labs, val_labs, train_s, val_s = train_test_split(
-        data, labs, sel_coeffs, stratify=labs, test_size=0.3
+    (
+        train_data,
+        val_data,
+        train_labs,
+        val_labs,
+        train_s,
+        val_s,
+        train_reps,
+        val_reps,
+        train_offsets,
+        val_offsets,
+    ) = train_test_split(
+        data, labs, sel_coeffs, reps, offsets, stratify=labs, test_size=0.3
     )
-    val_data, test_data, val_labs, test_labs, val_s, test_s = train_test_split(
-        val_data, val_labs, val_s, stratify=val_labs, test_size=0.5
+    (
+        val_data,
+        test_data,
+        val_labs,
+        test_labs,
+        val_s,
+        test_s,
+        val_reps,
+        test_reps,
+        val_offsets,
+        test_offsets,
+    ) = train_test_split(
+        val_data,
+        val_labs,
+        val_s,
+        val_reps,
+        val_offsets,
+        stratify=val_labs,
+        test_size=0.5,
     )
     return (
         train_data,
@@ -100,56 +148,12 @@ def split_partitions(data, labs, sel_coeffs):
         train_s,
         val_s,
         test_s,
-    )
-
-
-def filter_s_vals(
-    cls_idx,
-    ts_train_data,
-    ts_val_data,
-    ts_test_data,
-    train_labs,
-    train_s,
-    val_labs,
-    val_s,
-    test_labs,
-    test_s,
-):
-    cleaned_train_idxs_1 = np.where(train_labs[:, 0] == 0)
-    cleaned_train_idxs_2 = np.where(train_labs[:, cls_idx] == 1)
-    cleaned_train_idxs_3 = np.where(train_s > 0.0)
-    _cleaned_train_idxs = np.intersect1d(cleaned_train_idxs_1, cleaned_train_idxs_2)
-    cleaned_train_idxs = np.intersect1d(_cleaned_train_idxs, cleaned_train_idxs_3)
-
-    cleaned_val_idxs_1 = np.where(val_labs[:, 0] == 0)
-    cleaned_val_idxs_2 = np.where(val_labs[:, cls_idx] == 1)
-    cleaned_val_idxs_3 = np.where(val_s > 0.0)
-    _cleaned_val_idxs = np.intersect1d(cleaned_val_idxs_1, cleaned_val_idxs_2)
-    cleaned_val_idxs = np.intersect1d(_cleaned_val_idxs, cleaned_val_idxs_3)
-
-    cleaned_test_idxs_1 = np.where(test_labs[:, 0] == 0)
-    cleaned_test_idxs_2 = np.where(test_labs[:, cls_idx] == 1)
-    cleaned_test_idxs_3 = np.where(test_s > 0.0)
-    _cleaned_test_idxs = np.intersect1d(cleaned_test_idxs_1, cleaned_test_idxs_2)
-    cleaned_test_idxs = np.intersect1d(_cleaned_test_idxs, cleaned_test_idxs_3)
-
-    clean_train_svals = train_s[cleaned_train_idxs, :]
-    clean_val_svals = val_s[cleaned_val_idxs]
-    clean_test_svals = test_s[cleaned_test_idxs]
-
-    clean_train_data = ts_train_data[cleaned_train_idxs, :]
-    clean_val_data = ts_val_data[cleaned_val_idxs, :]
-    clean_test_data = ts_test_data[cleaned_test_idxs, :]
-    clean_test_labs = test_labs[cleaned_test_idxs, :]
-
-    return (
-        clean_train_data,
-        clean_val_data,
-        clean_test_data,
-        clean_test_labs,
-        clean_train_svals,
-        clean_val_svals,
-        clean_test_svals,
+        train_reps,
+        val_reps,
+        test_reps,
+        train_offsets,
+        val_offsets,
+        test_offsets,
     )
 
 
@@ -165,7 +169,6 @@ def fit_class_model(
 ):
     """
     Fits a given model using training/validation data, plots history after done.
-
     Args:
         out_dir (str): Base directory where data is located, model will be saved here.
         model (Model): Compiled Keras model.
@@ -175,7 +178,6 @@ def fit_class_model(
         val_data (np.arr): Validation data.
         val_labs (list[int]): OHE labels for validation set.
         experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
-
     Returns:
         Model: Fitted Keras model, ready to be used for accuracy characterization.
     """
@@ -200,7 +202,7 @@ def fit_class_model(
 
     earlystop = EarlyStopping(
         monitor=monitor,
-        min_delta=0.05,
+        min_delta=0.1,
         patience=20,
         verbose=1,
         mode="auto",
@@ -212,7 +214,7 @@ def fit_class_model(
     history = model.fit(
         x=train_data,
         y=train_target,
-        epochs=40,
+        epochs=100,
         verbose=2,
         callbacks=callbacks_list,
         validation_data=(val_data, val_target),
@@ -276,8 +278,8 @@ def fit_reg_model(
 
     earlystop = EarlyStopping(
         monitor=monitor,
-        min_delta=0.05,
-        patience=20,
+        min_delta=0.01,
+        patience=40,
         verbose=1,
         mode="auto",
         restore_best_weights=True,
@@ -288,7 +290,7 @@ def fit_reg_model(
     history = model.fit(
         x=train_data,
         y=train_target,
-        epochs=40,
+        epochs=100,
         verbose=2,
         callbacks=callbacks_list,
         validation_data=(val_data, val_target),
@@ -305,7 +307,7 @@ def fit_reg_model(
     save_model(
         model,
         os.path.join(
-            out_dir, "trained_models", f"{experiment_name}_{model.name}_{data_type}"
+            out_dir, "trained_models", f"REG_{experiment_name}_{model.name}_{data_type}"
         ),
     )
 
@@ -316,7 +318,9 @@ def evaluate_reg_model(
     model,
     test_data,
     test_labs,
-    test_s,
+    trans_test_s,
+    test_reps,
+    test_offsets,
     s_scaler,
     out_dir,
     scenarios,
@@ -337,21 +341,26 @@ def evaluate_reg_model(
         experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
         data_type (str): Whether data is aft or hfs.
     """
-
-    pred_s = model.predict(test_data)
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
+    trans_pred_s = model.predict(test_data)
     trues = np.argmax(test_labs, axis=1)
 
     if "log" in experiment_name:
-        pred_s = 10 ** (-pred_s)
-        test_s = 10 ** (-test_s)
+        pred_s = 10 ** (-trans_pred_s)
+        test_s = 10 ** (-trans_test_s)
     elif "minmax" in experiment_name:
-        pred_s = s_scaler.inverse_transform(pred_s.reshape(-1, 1))
-        test_s = s_scaler.inverse_transform(test_s.reshape(-1, 1))
+        pred_s = s_scaler.inverse_transform(trans_pred_s.reshape(-1, 1))
+        test_s = s_scaler.inverse_transform(trans_test_s.reshape(-1, 1))
+    else:
+        pred_s = trans_pred_s
+        test_s = trans_test_s
 
     pred_dict = {
-        "class": trues,
-        "true_sel_coeff": list(test_s),
-        "pred_sel_coeff": list(pred_s),
+        "rep": test_reps,
+        "offset": test_offsets.flatten(),
+        "class": [str_lab_dict[i] for i in trues],
+        "true_sel_coeff": test_s.flatten(),
+        "pred_sel_coeff": pred_s.flatten(),
     }
 
     pred_df = pd.DataFrame(pred_dict)
@@ -367,8 +376,8 @@ def evaluate_reg_model(
         index=False,
     )
 
-    print(
-        f"Mean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
+    logger.info(
+        f"\nMean absolute error for Sel Coeff predictions: {mean_absolute_error(test_s, pred_s)}"
     )
     pu.plot_sel_coeff_preds(
         trues,
@@ -387,6 +396,8 @@ def evaluate_class_model(
     model,
     test_data,
     test_labs,
+    test_reps,
+    test_offsets,
     out_dir,
     scenarios,
     experiment_name,
@@ -406,6 +417,7 @@ def evaluate_class_model(
         experiment_name (str): Descriptor of the sampling strategy used to generate the data. Used to ID the output.
         data_type (str): Whether data is aft or hfs.
     """
+    str_lab_dict = {value: key for key, value in lab_dict.items()}
 
     class_probs = model.predict(test_data)
     class_predictions = np.argmax(class_probs, axis=1)
@@ -417,8 +429,10 @@ def evaluate_class_model(
     pr_trues = np.array(list(trues))
 
     pred_dict = {
-        "true": trues,
-        "pred": class_predictions,
+        "rep": test_reps,
+        "offset": test_offsets.flatten(),
+        "true": [str_lab_dict[i] for i in trues],
+        "pred": [str_lab_dict[i] for i in class_predictions],
     }
     for str_lab in lab_dict:
         pred_dict[f"{str_lab}_scores"] = class_probs[:, lab_dict[str_lab]]
@@ -481,190 +495,214 @@ def evaluate_class_model(
 def main(ua):
     yaml_data = read_config(ua.yaml_file)
     work_dir = yaml_data["work dir"]
+    os.makedirs(os.path.join(work_dir, "images"), exist_ok=True)
 
     # Collect all the data
-    logger.info("Starting training process.")
-    if ua.hft:
-        type_list = ["aft", "hft"]
+    logger.info(
+        "Starting training process, shoulders are included in the training set."
+    )
+
+    ids, raw_reps, raw_ts_data, sweep_types, raw_sel_coeffs, raw_offsets = get_data(
+        ua.training_data, ua.data_type, yaml_data["win_size"]
+    )
+    lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
+
+    # Convert to numerical ohe IDs
+    num_ids = np.array([lab_dict[lab] for lab in ids])
+    raw_ohe_ids = to_categorical(
+        np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+    )
+
+    if ua.subsample_amount:
+        subsample_amount = ua.subsample_amount * len(set(ids))
+        # Subsample to test for training size effects
+        (
+            ts_data,
+            _,
+            ohe_ids,
+            _,
+            sel_coeffs,
+            _,
+            reps,
+            _,
+            offsets,
+            _,
+        ) = train_test_split(
+            raw_ts_data,
+            raw_ohe_ids,
+            raw_sel_coeffs,
+            raw_reps,
+            raw_offsets,
+            train_size=subsample_amount,
+            stratify=raw_ohe_ids,
+        )
     else:
-        type_list = ["aft"]
-    for data_type in type_list:
-        ids, ts_data, sweep_types, sel_coeffs = get_data(ua.training_data, data_type)
-        lab_dict = {str_id: int_id for int_id, str_id in enumerate(sweep_types)}
-
-        # Convert to numerical ohe IDs
-        num_ids = np.array([lab_dict[lab] for lab in ids])
-        ohe_ids = to_categorical(
-            np.array([lab_dict[lab] for lab in ids]), len(set(ids))
+        ts_data = raw_ts_data
+        ohe_ids = raw_ohe_ids
+        reps = raw_reps
+        sel_coeffs = raw_sel_coeffs
+        offsets = raw_offsets
+
+    logger.info(f"Data is subsampled to {len(ts_data)}")
+
+    class_weights = dict(
+        enumerate(
+            compute_class_weight(
+                class_weight="balanced", classes=np.unique(num_ids), y=num_ids
+            )
         )
+    )
+    logger.info(f"Class weights: {class_weights}")
 
-        class_weights = dict(
-            enumerate(
-                compute_class_weight(
-                    class_weight="balanced", classes=np.unique(num_ids), y=num_ids
-                )
-            )
+    if ua.data_type == "aft":
+        # Needs to be in correct dims order for Conv1D layer
+        datadim = ts_data.shape[1:]
+        logger.info(
+            f"{ua.data_type.upper()} TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
+        )
+    else:
+        logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
+        datadim = ts_data.shape[1:]
+        logger.info(f"{len(ts_data)} samples in dataset.")
+
+    logger.info("Splitting Partitions for Classification Task")
+    (
+        ts_train_data,
+        ts_val_data,
+        ts_test_data,
+        train_labs,
+        val_labs,
+        test_labs,
+        train_s,
+        val_s,
+        test_s,
+        _train_reps,
+        _val_reps,
+        test_reps,
+        train_offsets,
+        val_offsets,
+        test_offsets,
+    ) = split_partitions(ts_data, ohe_ids, sel_coeffs, reps, offsets)
+
+    # Time-series model training and evaluation
+    logger.info("Training time-series model.")
+
+    # Lazy switch for testing
+    model_type = "1dcnn"
+    if model_type == "1dcnn":
+        class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
+        reg_model = models.create_TS_reg_model(datadim)  # type: ignore
+    elif model_type == "transformer":
+        class_model = models.create_transformer_class_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
+            n_class=len(lab_dict),
+        )
+        reg_model = models.create_transformer_reg_model(
+            input_shape=ts_train_data.shape[1:],
+            head_size=256,
+            num_heads=4,
+            ff_dim=4,
+            num_transformer_blocks=4,
+            mlp_units=[128],
+            mlp_dropout=0.4,
+            dropout=0.25,
         )
-        print(f"Class weights: {class_weights}")
+    else:
+        logger.error("Need a model")
+        sys.exit(1)
 
-        if data_type == "aft":
-            # Needs to be in correct dims order for Conv1D layer
-            datadim = ts_data.shape[1:]
-            logger.info(
-                f"TS Data shape (samples, timepoints, alleles): {ts_data.shape}"
-            )
-        else:
-            logger.info(f"TS Data shape (samples, timepoints, haps): {ts_data.shape}")
-            datadim = ts_data.shape[1:]
-            logger.info(f"{len(ts_data)} samples in dataset.")
+    run_modes = ["class", "reg"]
+    if "class" in run_modes:
+        logger.info(f"\nRunning classification model for {ua.data_type}")
 
-        logger.info("Splitting Partition")
-        (
+        trained_class_model = fit_class_model(
+            work_dir,
+            class_model,
+            ua.data_type,
             ts_train_data,
-            ts_val_data,
-            ts_test_data,
             train_labs,
+            ts_val_data,
             val_labs,
+            ua.experiment_name,
+        )
+        evaluate_class_model(
+            trained_class_model,
+            ts_test_data,
             test_labs,
-            train_s,
-            val_s,
-            test_s,
-        ) = split_partitions(ts_data, ohe_ids, sel_coeffs)
-
-        # Time-series model training and evaluation
-        logger.info("Training time-series model.")
-
-        # Lazy switch for testing
-        model_type = "1dcnn"
-        if model_type == "1dcnn":
-            class_model = models.create_TS_class_model(datadim, len(lab_dict))  # type: ignore
-            reg_model = models.create_TS_reg_model(datadim)  # type: ignore
-        elif model_type == "transformer":
-            class_model = models.create_transformer_class_model(
-                input_shape=ts_train_data.shape[1:],
-                head_size=256,
-                num_heads=4,
-                ff_dim=4,
-                num_transformer_blocks=4,
-                mlp_units=[128],
-                mlp_dropout=0.4,
-                dropout=0.25,
-                n_class=len(lab_dict),
-            )
-            reg_model = models.create_transformer_reg_model(
-                input_shape=ts_train_data.shape[1:],
-                head_size=256,
-                num_heads=4,
-                ff_dim=4,
-                num_transformer_blocks=4,
-                mlp_units=[128],
-                mlp_dropout=0.4,
-                dropout=0.25,
+            test_reps,
+            test_offsets,
+            work_dir,
+            yaml_data["scenarios"],
+            ua.experiment_name,
+            ua.data_type,
+            lab_dict,
+        )
+
+    if "reg" in run_modes:
+        for idx, scenario in enumerate(yaml_data["scenarios"][1:], start=1):
+            # print(reg_model.summary())
+
+            mode = "minmax"
+            if mode == "log":
+                mm_scaler = None
+                trvals = -np.log10(train_s)
+                vvals = -np.log10(val_s)
+                tevals = -np.log10(test_s)
+            elif mode == "minmax":
+                mm_scaler = scale_sel_coeffs(train_s)
+                trvals = mm_scaler.transform(train_s)
+                vvals = mm_scaler.transform(val_s)
+                tevals = mm_scaler.transform(test_s)
+            else:
+                mm_scaler = None
+                trvals = train_s
+                vvals = val_s
+                tevals = test_s
+
+            logger.info(
+                f"{ua.data_type.upper()} Regression {scenario.upper()} TS Data shape (samples, timepoints, alleles/haplotypes): {ts_train_data.shape} {ts_val_data.shape} {ts_test_data.shape}"
             )
-        else:
-            print("Need a model")
-            sys.exit(1)
 
-        run_modes = ["reg"]  # ["class", "reg"]
-        if "class" in run_modes:
-            print(class_model.summary())
+            plot = True
+            if plot:
+                pu.plot_s_vs_freqs(
+                    train_s,
+                    ts_train_data[:, -1, int(ts_train_data.shape[-1] / 2)]
+                    - ts_train_data[:, 0, int(ts_train_data.shape[-1] / 2)],
+                    scenario,
+                    work_dir,
+                    ua.experiment_name,
+                    mode,
+                )
 
-            trained_class_model = fit_class_model(
+            trained_reg_model = fit_reg_model(
                 work_dir,
-                class_model,
-                data_type,
+                reg_model,
+                ua.data_type,
                 ts_train_data,
-                train_labs,
+                trvals,
                 ts_val_data,
-                val_labs,
-                ua.experiment_name,
+                vvals,
+                ua.experiment_name + f"_{scenario}_{mode}",
             )
-            evaluate_class_model(
-                trained_class_model,
+            evaluate_reg_model(
+                trained_reg_model,
                 ts_test_data,
                 test_labs,
+                tevals,
+                test_reps,
+                test_offsets,
+                mm_scaler,
                 work_dir,
                 yaml_data["scenarios"],
-                ua.experiment_name,
-                data_type,
+                ua.experiment_name + f"_{scenario}_{mode}",
+                ua.data_type,
                 lab_dict,
             )
-        if "reg" in run_modes:
-            # Remove the "neutral" scenario
-            # Create separate datasets for sdn/ssv regression models
-            for idx, scenario in zip([1, 2], ["sdn", "ssv"]):
-                (
-                    clean_train_data,
-                    clean_val_data,
-                    clean_test_data,
-                    clean_test_labs,
-                    clean_train_svals,
-                    clean_val_svals,
-                    clean_test_svals,
-                ) = filter_s_vals(
-                    idx,
-                    ts_train_data,
-                    ts_val_data,
-                    ts_test_data,
-                    train_labs,
-                    train_s,
-                    val_labs,
-                    val_s,
-                    test_labs,
-                    test_s,
-                )
-
-                # print(reg_model.summary())
-                logger.info(
-                    f"TS Data shape (samples, timepoints, alleles): {clean_train_data.shape}"
-                )
-
-                mm_scaler = scale_sel_coeffs(clean_train_svals)
-
-                mode = "log"
-                if mode == "log":
-                    trvals = -np.log10(clean_train_svals)
-                    vvals = -np.log10(clean_val_svals)
-                    tevals = -np.log10(clean_test_svals)
-                elif mode == "minmax":
-                    trvals = mm_scaler.transform(clean_train_svals)
-                    vvals = mm_scaler.transform(clean_val_svals)
-                    tevals = mm_scaler.transform(clean_test_svals)
-                else:
-                    trvals = clean_train_svals
-                    vvals = clean_val_svals
-                    tevals = clean_test_svals
-
-                plot = True
-                if plot:
-                    pu.plot_s_vs_freqs(
-                        clean_train_svals,
-                        clean_train_data[:, -1, 25] - clean_train_data[:, 0, 25],
-                        scenario,
-                        work_dir,
-                        ua.experiment_name,
-                        mode,
-                    )
-
-                trained_reg_model = fit_reg_model(
-                    work_dir,
-                    reg_model,
-                    data_type,
-                    clean_train_data,
-                    trvals,
-                    clean_val_data,
-                    vvals,
-                    ua.experiment_name + f"_{scenario}_{mode}",
-                )
-                evaluate_reg_model(
-                    trained_reg_model,
-                    clean_test_data,
-                    clean_test_labs,
-                    tevals,
-                    mm_scaler,
-                    work_dir,
-                    yaml_data["scenarios"],
-                    ua.experiment_name + f"_{scenario}_{mode}",
-                    data_type,
-                    lab_dict,
-                )
diff --git a/timesweeper/utils/gen_utils.py b/timesweeper/utils/gen_utils.py
index d4dbcf8..60e7a0d 100644
--- a/timesweeper/utils/gen_utils.py
+++ b/timesweeper/utils/gen_utils.py
@@ -36,100 +36,6 @@ def get_rep_id(filepath):
             continue
 
 
-def write_fit(fit_dict, outfile, benchmark):
-    """
-    Writes FIT predictions to file.
-
-    Args:
-        fit_dict (dict): FIT p values and SNP information.
-        outfile (str): File to write results to.
-    """
-    inv_pval = [1 - i[1] for i in fit_dict.values()]
-    if benchmark:
-        chroms, bps, mut_type = zip(*fit_dict.keys())
-        predictions = pd.DataFrame(
-            {"Chrom": chroms, "BP": bps, "Mut_Type": mut_type, "Inv_pval": inv_pval}
-        )
-    else:
-        chroms, bps = zip(*fit_dict.keys())
-        predictions = pd.DataFrame({"Chrom": chroms, "BP": bps, "Inv_pval": inv_pval})
-
-    predictions.dropna(inplace=True)
-    predictions.sort_values(["Chrom", "BP"], inplace=True)
-    predictions.to_csv(
-        os.path.join(outfile), header=True, index=False, sep="\t",
-    )
-
-
-def write_preds(results_dict, outfile, benchmark):
-    """
-    Writes NN predictions to file.
-
-    Args:
-        results_dict (dict): SNP NN prediction scores and window edges.
-        outfile (str): File to write results to.
-    """
-    lab_dict = {0: "Neut", 1: "SSV", 2: "SDN"}
-    if benchmark:
-        chroms, bps, mut_type, true_sel_coeff = zip(*results_dict.keys())
-    else:
-        chroms, bps = zip(*results_dict.keys())
-
-    neut_scores = [i[0][0] for i in results_dict.values()]
-    sdn_scores = [i[0][1] for i in results_dict.values()]
-    ssv_scores = [i[0][2] for i in results_dict.values()]
-    sel_preds = [i[1] for i in results_dict.values()]
-    left_edges = [i[2] for i in results_dict.values()]
-    right_edges = [i[3] for i in results_dict.values()]
-    classes = [lab_dict[np.argmax(i[0])] for i in results_dict.values()]
-
-    if benchmark:
-        predictions = pd.DataFrame(
-            {
-                "Chrom": chroms,
-                "BP": bps,
-                "Mut_Type": mut_type,
-                "Class": classes,
-                "Sweep_Score": [i + j for i, j in zip(sdn_scores, ssv_scores)],
-                "True_Sel_Coeff": true_sel_coeff,
-                "Sel_Coeff": sel_preds,
-                "Neut_Score": neut_scores,
-                "SDN_Score": sdn_scores,
-                "SSV_Score": ssv_scores,
-                "Win_Start": left_edges,
-                "Win_End": right_edges,
-            }
-        )
-    else:
-        predictions = pd.DataFrame(
-            {
-                "Chrom": chroms,
-                "BP": bps,
-                "Class": classes,
-                "Sweep_Score": [i + j for i, j in zip(sdn_scores, ssv_scores)],
-                "Sel_Coeff": sel_preds,
-                "Neut_Score": neut_scores,
-                "SDN_Score": sdn_scores,
-                "SSV_Score": ssv_scores,
-                "Win_Start": left_edges,
-                "Win_End": right_edges,
-            }
-        )
-        predictions = predictions[predictions["Neut_Score"] < 0.5]
-
-    predictions.sort_values(["Chrom", "BP"], inplace=True)
-
-    if not os.path.exists(outfile):
-        predictions.to_csv(outfile, header=True, index=False, sep="\t")
-    else:
-        predictions.to_csv(outfile, mode="a", header=False, index=False, sep="\t")
-
-    bed_df = predictions[["Chrom", "Win_Start", "Win_End", "BP"]]
-    bed_df.to_csv(
-        outfile.replace(".csv", ".bed"), mode="a", header=False, index=False, sep="\t"
-    )
-
-
 def get_logger(module_name):
     logging.basicConfig()
     logger = logging.getLogger(module_name)
diff --git a/timesweeper/utils/snp_utils.py b/timesweeper/utils/snp_utils.py
index d571c67..fc333e4 100644
--- a/timesweeper/utils/snp_utils.py
+++ b/timesweeper/utils/snp_utils.py
@@ -1,12 +1,12 @@
 import allel
 import numpy as np
 
+
 # General util functions
-def read_vcf(vcf_file, benchmark):
+def read_vcf(vcf_file, samples_list, benchmark):
     """
     Loads VCF file and grabs relevant fields.
     For generating training data from simulated VCFs, which are typically small.
-
     Args:
         vcf_file (str): Path to vcf file.
         benchmark (bool): Whether to look for Mut_Type or not.
@@ -23,9 +23,17 @@ def read_vcf(vcf_file, benchmark):
         ]
     else:
         fields = ["variants/CHROM", "variants/POS", "calldata/GT"]
-
-    vcf = allel.read_vcf(vcf_file, fields=fields)
-
+    if samples_list:
+        vcf = allel.read_vcf(
+            vcf_file,
+            fields=fields,
+            samples=samples_list,
+            numbers={"variants/MT": 100, "variants/S": 100},
+        )
+    else:
+        vcf = allel.read_vcf(
+            vcf_file, fields=fields, numbers={"variants/MT": 100, "variants/S": 100}
+        )
     return vcf
 
 
@@ -53,7 +61,10 @@ def get_vcf_iter(vcf_file, benchmark):
         fields = ["variants/CHROM", "variants/POS", "calldata/GT"]
 
     _fields, _samples, _headers, vcf_iter = allel.iter_vcf_chunks(
-        vcf_file, fields=fields, chunk_length=100000
+        vcf_file,
+        fields=fields,
+        chunk_length=100000,
+        numbers={"variants/MT": 20, "variants/S": 20},
     )
 
     return vcf_iter
@@ -84,14 +95,10 @@ def make_loc_tups(vcf, benchmark):
         list[tuple]: List of tuples with (chrom, pos, mut).
     """
     if benchmark:
-        return list(
-            zip(
-                vcf["variants/CHROM"],
-                vcf["variants/POS"],
-                vcf["variants/MT"],
-                vcf["variants/S"],
-            )
-        )
+        MTs = np.nanmax(vcf["variants/MT"], axis=1).flatten()
+        Ss = np.nanmax(vcf["variants/S"], axis=1).flatten()
+
+        return list(zip(vcf["variants/CHROM"], vcf["variants/POS"], MTs, Ss,))
     else:
         return list(zip(vcf["variants/CHROM"], vcf["variants/POS"]))
 
@@ -139,7 +146,7 @@ def split_arr(arr, samp_sizes):
 
     Args:
         arr (np.arr): SNP or Haplotype array with all timepoints in flat structure.
-        samp_sizes (list(int)): List of chromosomes (not individuals) to index from the array.
+        samp_sizes (list(int)): List of individuals to index from the array.
 
     Returns:
         list[np.arr]: Time-serialized list of arrays of SNP or haplotype data.
@@ -164,12 +171,24 @@ def get_vel_minor_alleles(ts_genos, max_allele):
     Returns:
         np.arr: Array of indices of minor alleles.
     """
-    # Shape is (snps, counts)
-    # Highest velocity allele wins
-    last_genos = allel.GenotypeArray(ts_genos[-1]).count_alleles(max_allele=max_allele)
-    first_genos = allel.GenotypeArray(ts_genos[0]).count_alleles(max_allele=max_allele)
+    if len(ts_genos) == 1:
+        # Single timepoint
+        last_genos = allel.GenotypeArray(ts_genos[0]).count_alleles(
+            max_allele=max_allele
+        )
+        return np.argmax(last_genos, axis=1), None, None
+
+    else:
+        # Shape is (snps, counts)
+        # Highest velocity allele wins
+        last_genos = allel.GenotypeArray(ts_genos[-1]).count_alleles(
+            max_allele=max_allele
+        )
+        first_genos = allel.GenotypeArray(ts_genos[0]).count_alleles(
+            max_allele=max_allele
+        )
 
-    return np.argmax(last_genos - first_genos, axis=1)
+        return np.argmax(last_genos - first_genos, axis=1), first_genos, last_genos
 
 
 def get_last_minor_alleles(ts_genos, max_allele):
@@ -186,7 +205,7 @@ def get_last_minor_alleles(ts_genos, max_allele):
     # Shape is (snps, counts)
     last_genos = allel.GenotypeArray(ts_genos[-1]).count_alleles(max_allele=max_allele)
 
-    return np.argmax(last_genos, axis=1)
+    return np.argmax(last_genos, axis=1), None, None
 
 
 def calc_maft(snp, min_allele_idx):
@@ -200,7 +219,11 @@ def calc_maft(snp, min_allele_idx):
     Returns:
         float: Minor allele frequency (MAF) at a given timepoint.
     """
-    return np.divide(snp[min_allele_idx], snp.sum())
+    try:
+        maf = np.divide(snp[min_allele_idx], snp.sum())
+        return maf
+    except RuntimeWarning:
+        return 0.0
 
 
 def get_allele_counts(snp, min_allele_idx):


Test patch:
diff --git a/timesweeper/tests/test_find_sweeps.py b/timesweeper/tests/test_find_sweeps.py
index 9ad534c..cbba8f3 100644
--- a/timesweeper/tests/test_find_sweeps.py
+++ b/timesweeper/tests/test_find_sweeps.py
@@ -1,11 +1,11 @@
 import numpy as np
-from ..utils.gen_utils import (
+from timesweeper.utils.gen_utils import (
     get_scenario_from_filename,
     get_rep_id,
     add_file_label,
     read_config,
 )
-from ..find_sweeps_vcf import get_window_idxs
+from timesweeper.find_sweeps_vcf import get_window_idxs
 import pytest
 
 scenarios = ["sdn", "ssv" "neutral"]
diff --git a/timesweeper/tests/test_fit.py b/timesweeper/tests/test_fit.py
index b73c867..5233af2 100644
--- a/timesweeper/tests/test_fit.py
+++ b/timesweeper/tests/test_fit.py
@@ -1,4 +1,4 @@
-from ..utils import frequency_increment_test as fit
+from timesweeper.utils import frequency_increment_test as fit
 import numpy as np
 
 
diff --git a/timesweeper/tests/test_make_training_features.py b/timesweeper/tests/test_make_training_features.py
index ca1e5fd..8227576 100644
--- a/timesweeper/tests/test_make_training_features.py
+++ b/timesweeper/tests/test_make_training_features.py
@@ -1,4 +1,4 @@
-from .. import make_training_features as mtf
+from timesweeper. import make_training_features as mtf
 import numpy as np
 
 
diff --git a/timesweeper/utils/fishers_exact_test.py b/timesweeper/utils/fishers_exact_test.py
deleted file mode 100644
index 3038049..0000000
--- a/timesweeper/utils/fishers_exact_test.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from scipy import stats
-import numpy as np
-import snp_utils as su
-
-
-def fet(maj_counts, min_counts):
-    """
-    Run Fisher's exact test using minor allele frequency vector given by prep_ts_aft.add()
-
-    Args:
-        maf_vector (list[float]): Time-series of allele frequencies for highest velocity allele in a given SNP.add()
-
-    Returns:
-        float: p-value, result of FET on first and last minor/major allele frequency contingency table.
-    """
-    return stats.fisher_exact(
-        [[maj_counts[0], maj_counts[-1]], [min_counts[0], min_counts[-1]],]
-    )[1]
-
-
-def test_fet():
-    maj_counts = np.random.randint(0, 20, (20))
-    min_counts = np.random.randint(0, 5, (20))
-    fish_res = fet(maj_counts, min_counts)
-    print(fish_res)
-
-
-def main():
-    test_fet()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/timesweeper/utils/frequency_increment_test.py b/timesweeper/utils/frequency_increment_test.py
deleted file mode 100644
index a29dc70..0000000
--- a/timesweeper/utils/frequency_increment_test.py
+++ /dev/null
@@ -1,56 +0,0 @@
-import warnings
-
-import pandas as pd
-from scipy.stats import ttest_1samp
-
-warnings.filterwarnings("ignore", category=RuntimeWarning)
-# https://www.genetics.org/content/196/2/509
-
-df = pd.DataFrame
-
-
-def getRescaledIncs(freqs, gens):
-    """
-    Calculates allele increments for the FIT.
-
-    Args:
-        freqs (List[int]): Frequency of a given allele at each site.
-        gens (List[int]): Generations that each allele is sampled.
-
-    Returns:
-        List[float]: Rescaled frequency increments that can be used in t-test.
-    """
-    incs = []
-    i = 0
-    # advance to first non-zero freq
-    while i < len(freqs) and (freqs[i] == 0 or freqs[i] == 1):
-        i += 1
-    if i < len(freqs):
-        prevFreq = freqs[i]
-        prevGen = gens[i]
-        i += 1
-        while i < len(freqs):
-            if freqs[i] != 0 and freqs[i] != 1:
-                num = freqs[i] - prevFreq
-                denom = ((2 * prevFreq) * (1 - prevFreq) * (gens[i] - prevGen)) ** 0.5
-                incs.append(num / denom)
-                prevFreq = freqs[i]
-                prevGen = gens[i]
-            i += 1
-
-    return incs
-
-
-def fit(freqs, gens):
-    """
-    Calculate FIT by performing 1-Sided Student t-test on frequency increments.
-
-    Args:
-        freqs (List[int]): Frequencies at all given generations for targeted alleles.
-        gens (List[int]): Generations sampled.
-
-    Returns:
-        List[int]: t-test results.
-    """
-    rescIncs = getRescaledIncs(freqs, gens)
-    return ttest_1samp(rescIncs, 0)