snputils.visualization.admixture_viz

  1import subprocess
  2from typing import List, Optional, Dict, Any
  3import os
  4import sys
  5import argparse
  6import shutil
  7
  8
  9class PongVizError(Exception):
 10    """Custom exception for Pong visualization errors"""
 11    pass
 12
 13
 14class FileMapError(Exception):
 15    """Custom exception for filemap creation errors"""
 16    pass
 17
 18
 19def pong_viz(folder_runs: str, output_dir: str, k: Optional[int] = None, min_k: Optional[int] = None, max_k: Optional[int] = None,
 20             runs: List[int] = None, run_prefix: str = 'train', ind2pop_path: Optional[str] = None,
 21             pop_names_path: Optional[str] = None, color_list_path: Optional[str] = None, verbose: bool = False) -> None:
 22    """
 23    Executes Pong visualization with the specified parameters.
 24    """
 25
 26    # Check if pong is installed
 27    if not shutil.which('pong'):  # TODO: replace pong dependency with a custom implementation
 28        raise PongVizError("Pong is not installed. Please install it using 'pip install pong'")
 29
 30    # Validate input parameters
 31    if not os.path.isdir(folder_runs):
 32        raise PongVizError(f"Input folder does not exist: {folder_runs}")
 33    if not output_dir:
 34        raise PongVizError("Output directory must be specified")
 35    os.makedirs(output_dir, exist_ok=True)
 36
 37    # Optional files validation
 38    optional_files = {'ind2pop_path': ind2pop_path, 'pop_names_path': pop_names_path, 'color_list_path': color_list_path}
 39    for name, path in optional_files.items():
 40        if path and not os.path.isfile(path):
 41            raise PongVizError(f"Specified {name} file does not exist: {path}")
 42
 43    try:
 44        filemap_path = create_filemap(folder_runs, k, min_k, max_k, runs, run_prefix)
 45
 46        if verbose:
 47            print(f"Created filemap at: {filemap_path}")
 48
 49        # Build command
 50        cmd = ["pong", "-m", filemap_path, "-o", output_dir]
 51        if ind2pop_path:
 52            cmd.extend(["-i", ind2pop_path])
 53        if pop_names_path:
 54            cmd.extend(["-n", pop_names_path])
 55        if color_list_path:
 56            cmd.extend(["-l", color_list_path])
 57
 58        if verbose:
 59            print("Executing command:", " ".join(cmd))
 60
 61        # Execute pong
 62        subprocess.run(cmd)
 63
 64        if verbose:
 65            print("Pong execution successful")
 66            print("Results saved in:", output_dir)
 67
 68    except FileMapError as e:
 69        raise PongVizError(f"Error creating filemap: {str(e)}")
 70    except subprocess.CalledProcessError as e:
 71        raise PongVizError(f"Error executing Pong: {str(e)}\nOutput: {e.output}")
 72    except Exception as e:
 73        raise PongVizError(f"Unexpected error: {str(e)}")
 74
 75
 76def create_filemap(folder: str, k: Optional[int] = None, min_k: Optional[int] = None, max_k: Optional[int] = None,
 77                   runs: List[int] = None, run_prefix: str = 'train_demo',) -> str:
 78    """
 79    Creates a filemap for training files organized by k values and runs and saves it to a file.
 80
 81    Args:
 82        folder (str): Base folder path
 83        k (Optional[int]): Single k value to process. If specified, min_k and max_k are ignored
 84        min_k (Optional[int]): Minimum k value for range processing
 85        max_k (Optional[int]): Maximum k value for range processing
 86        runs (List[int]): List of run numbers
 87        run_prefix (str): Prefix for the run files (default: 'train_demo')
 88
 89    Returns:
 90        str: Path to saved file
 91
 92    Raises:
 93        FileMapError: If invalid parameters are provided or if configuration is incorrect
 94    """
 95    # Input validation
 96    if not folder:
 97        raise FileMapError("Folder path must be specified")
 98    if not runs:
 99        raise FileMapError("Runs list must be provided and non-empty")
100    if not all(isinstance(run, int) and run > 0 for run in runs):
101        raise FileMapError("All runs must be positive integers")
102    if k is not None:
103        if k < 1:
104            raise FileMapError("k must be a positive integer")
105        k_values = [k]
106    else:
107        if min_k is None or max_k is None:
108            raise FileMapError("Either k or both min_k and max_k must be specified")
109        if min_k > max_k:
110            raise FileMapError("min_k cannot be greater than max_k")
111        if min_k < 1:
112            raise FileMapError("min_k must be a positive integer")
113        k_values = range(min_k, max_k + 1)
114
115    # Write file:
116    filemap = {}
117    try:
118        for i, k_val in enumerate(k_values):
119            for run_num in range(runs[i]):
120                key = f'k{k_val}r{run_num+1}'
121                if folder != '.':
122                    path = folder + '/' + f'run{run_num+1}' + '/' + f'{run_prefix}.{k_val}.Q'
123                else:
124                    path = f'run{run_num+1}' + '/' + f'{run_prefix}.{k_val}.Q'
125
126                filemap[key] = (k_val, path)
127
128        filemap_path = os.path.join('pong_filemap')
129
130        with open(filemap_path, 'w') as f:
131            for i, (key, (k_val, path)) in enumerate(filemap.items()):
132                f.write(f"{key}\t{k_val}\t{path}\n")
133
134                if k_val in k_values:
135                    index = k_values.index(k_val)
136                    if (i + 1) % runs[index] == 0:
137                        f.write("\n")
138
139        return filemap_path
140
141    except Exception as e:
142        raise FileMapError(f"Error creating filemap: {str(e)}")
143
144
145def parse_pong_args() -> Dict[str, Any]:
146    """
147    Parse command line arguments for pong visualization.
148
149    Returns:
150        Dict[str, Any]: Dictionary containing all parsed arguments ready for pong_viz
151
152    Raises:
153        SystemExit: If argument parsing fails or --help is used
154    """
155    parser = argparse.ArgumentParser(
156        description='Execute Pong visualization with specified parameters',
157        formatter_class=argparse.ArgumentDefaultsHelpFormatter
158    )
159
160    parser.add_argument('folder_runs', help='Base folder containing the run files')
161    parser.add_argument('output_dir', help='Directory where results will be saved')
162
163    k_group = parser.add_mutually_exclusive_group(required=True)
164    k_group.add_argument('-k', '--k', type=int, help='Single k value to process')
165    k_group.add_argument('--k-range', nargs=2, type=int, metavar=('MIN_K', 'MAX_K'), help='Range of k values (min_k max_k)')
166
167    parser.add_argument('--runs', type=int, nargs='+', help='Number of runs for each K')
168    parser.add_argument('--prefix', default='train_demo', help='Prefix for run files')
169    parser.add_argument('-i', '--ind2pop', help='Path to individual-to-population mapping file')
170    parser.add_argument('-n', '--pop-names', help='Path to population names file')
171    parser.add_argument('-c', '--color-list', help='Path to color list file')
172    parser.add_argument('-v', '--verbose', action='store_true', help='Print detailed information during execution')
173    args = parser.parse_args()
174
175    params = {'folder_runs': args.folder_runs,
176              'output_dir': args.output_dir,
177              'runs': args.runs,
178              'run_prefix': args.prefix,
179              'verbose': args.verbose}
180
181    if args.k is not None:
182        params['k'] = args.k
183    else:
184        params['min_k'] = args.k_range[0]
185        params['max_k'] = args.k_range[1]
186
187    if args.ind2pop:
188        params['ind2pop_path'] = args.ind2pop
189    if args.pop_names:
190        params['pop_names_path'] = args.pop_names
191    if args.color_list:
192        params['color_list_path'] = args.color_list
193
194    return params
195
196
197if __name__ == "__main__":
198    try:
199        params = parse_pong_args()
200        pong_viz(**params)
201
202    except (FileMapError, PongVizError) as e:
203        print(f"Error: {e}", file=sys.stderr)
204        sys.exit(1)
205    except Exception as e:
206        print(f"Unexpected error: {e}", file=sys.stderr)
207        sys.exit(1)