amrpredict 0.3.0__py3-none-any.whl

amrpredict/__init__.py

@@ -0,0 +1,259 @@
+"""amrpredict — antimicrobial resistance prediction from genomic data.
+
+The models the project's web app serves, behind one flat API:
+
+>>> import amrpredict
+>>> amrpredict.forecast('ciprofloxacin', taxon_id=562, mic_value=4)['prediction']
+'Susceptible'
+
+Trained models:
+
+* ``forecast`` — resistance from metadata and an MIC value (LightGBM).
+* ``predict_fasta`` — resistance from a complete genome: the gene model
+  (LightGBM on AMRFinderPlus genes) when AMRFinderPlus is installed, the k-mer
+  model (LightGBM on 4-mer frequencies) otherwise.
+* ``forecast_trend`` — the yearly % resistant of a genus and antibiotic, and its
+  forecast for the coming years (LightGBM on BV-BRC laboratory results).
+* ``plan_treatment`` — which drug to give each week, from an RL agent (PPO)
+  trained on the simulation, against fixed schedules.
+
+``simulate_timeline`` projects week-by-week resistance under treatment and is
+a biological simulation, not a trained model.
+"""
+from . import registry
+from ._paths import default_model_dir
+
+__version__ = '0.3.0'
+
+__all__ = [
+    'forecast',
+    'predict_fasta',
+    'identify_species',
+    'forecast_trend',
+    'trend_series',
+    'plan_treatment',
+    'simulate_timeline',
+    'antibiotics',
+    'status',
+    'registry',
+    'default_model_dir',
+    'LGBMResistancePredictor',
+    'GenomeModelPredictor',
+    'TimelineMLForecaster',
+    'RLAgent',
+    'MutationTimelinePredictor',
+    '__version__',
+]
+
+
+def forecast(antibiotic, *, taxon_id=None, mic_value=None, mic_sign=None,
+             genus='unknown', species='unknown', threshold=None, model_dir=None):
+    """Forecast resistance from genome metadata and (optionally) an MIC value.
+
+    Args:
+        antibiotic: Antibiotic name, case-insensitive (e.g. ``'ciprofloxacin'``).
+        taxon_id: NCBI taxonomy ID of the organism, e.g. ``562`` for *E. coli*.
+        mic_value: Minimum inhibitory concentration in mg/L.
+        mic_sign: Comparator recorded with the MIC — ``'='``, ``'>'``, ``'<='`` …
+        genus: Genus name, e.g. ``'Escherichia'``.
+        species: Species name, e.g. ``'coli'``.
+        threshold: Probability at or above which the call is ``'Resistant'``.
+            Default: the threshold the bundled model was promoted with
+            (``status()['lgbm_forecasting']['default_threshold']``).
+        model_dir: Override the bundled artifacts.
+
+    Returns:
+        dict with ``prediction``, ``probability`` (calibrated), ``confidence``,
+        ``drug_class``, ``model_run``, ``evidence`` (which inputs the model
+        recognised) and related fields.
+    """
+    return registry.lgbm(model_dir).predict(
+        antibiotic, taxon_id=taxon_id, mic_value=mic_value, mic_sign=mic_sign,
+        genus=genus, species=species, threshold=threshold,
+    )
+
+
+def predict_fasta(fasta_text, antibiotic, *, threshold=None, genus=None, species=None,
+                  use_genes=None, model_dir=None):
+    """Predict resistance from a complete genome assembly (at least 100,000 bp).
+
+    Args:
+        fasta_text: FASTA contents as a string; every contig is read.
+        antibiotic: Antibiotic name, case-insensitive.
+        threshold: Probability at or above which the call is ``'Resistant'``.
+            Default: the model's own threshold, chosen on validation genomes.
+        genus, species: The organism, if known (the k-mer model uses them).
+        use_genes: ``None`` (default) uses the gene model when AMRFinderPlus is
+            installed and the k-mer model otherwise; ``False`` always uses the
+            k-mer model; ``True`` requires the gene model.
+        model_dir: Override the bundled artifacts.
+
+    Returns:
+        dict with ``prediction``, ``probability``, ``confidence``,
+        ``gc_content``, ``sequence_length``, ``model_used``; ``top_kmers`` from
+        the k-mer model; ``genes_found`` and ``species_detected`` from the gene
+        model. Bad input gives a dict with ``error`` instead of raising.
+    """
+    if use_genes is False:
+        model = registry.genome_kmers(model_dir)
+    else:
+        model = registry.genome(model_dir)
+        if use_genes and not model.uses_genes:
+            return {'error': 'the gene model needs AMRFinderPlus (the `amrfinder` program), '
+                             'which was not found; pass use_genes=None to use the k-mer model'}
+    return model.predict(fasta_text, antibiotic, threshold, genus, species)
+
+
+def identify_species(fasta_text, *, model_dir=None):
+    """Species and genus of a complete genome from its 6-mer profile.
+
+    Returns:
+        dict with ``species``, ``genus`` (``None`` when no reference is close)
+        and ``distance``; or a dict with ``error``.
+    """
+    from .genome import GENE_FOLDER, GenomeModelPredictor, SpeciesProfiles, genome_counts
+    import os
+    path = os.path.join(model_dir or default_model_dir(), GENE_FOLDER, 'species_profiles.npz')
+    counts, _, length = genome_counts(fasta_text)
+    if length < GenomeModelPredictor.MIN_BP:
+        return {'error': f'genome too short ({length:,} bp): upload a complete assembly',
+                'sequence_length': length}
+    found = SpeciesProfiles(path).identify(counts) or {'species': None, 'genus': None, 'distance': None}
+    return {k: (str(v) if isinstance(v, str) else v) for k, v in found.items()}
+
+
+def forecast_trend(genus, antibiotic, *, years_ahead=5, model_dir=None):
+    """Yearly % resistant of a genus and antibiotic, and the timeline ML model's forecast.
+
+    Args:
+        genus: Genus name, case-insensitive (e.g. ``'Salmonella'``);
+            ``trend_series()`` lists those with enough history.
+        antibiotic: Antibiotic name, case-insensitive.
+        years_ahead: Years to forecast, 1 to 10. The 80% band widens with each.
+        model_dir: Override the bundled artifacts.
+
+    Returns:
+        dict with ``history`` (year, resistant_pct, n lab results),
+        ``forecast`` (year, resistant_pct, lo, hi), the model's held-out
+        ``metrics`` against the baselines and a ``label``; or a dict with
+        ``error`` when the pair has too little history.
+    """
+    years_ahead = int(years_ahead)
+    if not 1 <= years_ahead <= 10:
+        return {'error': 'years_ahead must be from 1 to 10'}
+    model = registry.trend(model_dir)
+    if not model.is_trained:
+        return {'error': 'the timeline ML model is not bundled'}
+    result = model.forecast(genus, antibiotic, years_ahead=years_ahead)
+    return result or {'error': f'{genus} has too little yearly lab history for {antibiotic} to forecast'}
+
+
+def trend_series(model_dir=None):
+    """``{genus: [antibiotic, ...]}`` that ``forecast_trend`` can forecast."""
+    return registry.trend(model_dir).series
+
+
+def plan_treatment(antibiotic, *, n_weeks=52, seed=42, model_dir=None):
+    """Which drug to give each week: the RL agent against fixed schedules.
+
+    The agent (PPO) was trained on the cycling simulation for ciprofloxacin,
+    gentamicin and imipenem; ``antibiotic`` is the first drug, the other two
+    are the alternatives. It learned from simulated episodes, not patients.
+
+    Args:
+        antibiotic: One of the agent's drugs, case-insensitive.
+        n_weeks: Weeks to plan, 1 to 104 (its training horizon).
+        seed: Seed of the simulated episode.
+        model_dir: Override the bundled artifacts.
+
+    Returns:
+        dict with ``drugs``, ``policies`` (the agent and each baseline: the drug
+        per week, resistance curves, ``failure_week``, ``total_reward``) and
+        ``best``; or a dict with ``error``.
+    """
+    from .amr_constants import normalize_antibiotic
+    agent = registry.rl(model_dir)
+    if agent is None:
+        return {'error': 'the RL agent is not bundled'}
+    n_weeks = int(n_weeks)
+    if not 1 <= n_weeks <= agent.train_weeks:
+        return {'error': f'n_weeks must be from 1 to {agent.train_weeks}'}
+    block = agent.block(normalize_antibiotic(antibiotic), n_weeks, seed=seed)
+    return block or {'error': f'the RL agent was trained on {", ".join(agent.drugs)} only'}
+
+
+def simulate_timeline(fasta_text, antibiotic, *, n_weeks=8, seed=42, model_dir=None):
+    """Simulate week-by-week resistance evolution under antibiotic pressure.
+
+    This is a biological simulation, not a trained model — it carries no AUC
+    or accuracy score. Treat its output as illustrative. The same inputs and
+    seed give the same output, and the three shares sum to 100 every week.
+    For a trained forecast of yearly resistance, use ``forecast_trend``.
+
+    Args:
+        fasta_text: FASTA contents as a string.
+        antibiotic: Antibiotic name, case-insensitive.
+        n_weeks: Number of weeks to project.
+        seed: Seed for the random parts (mutation counts and types).
+        model_dir: Override the bundled artifacts.
+
+    Returns:
+        dict with a ``weeks`` series, mutation events and gene activations.
+    """
+    return registry.timeline(model_dir).predict(fasta_text, antibiotic, n_weeks, seed=seed)
+
+
+def antibiotics(model='genome', model_dir=None):
+    """Sorted antibiotic names a model was trained on.
+
+    Args:
+        model: ``'genome'`` (``predict_fasta``), ``'forecast'`` (``forecast``),
+            ``'trend'`` (``forecast_trend``) or ``'rl'`` (``plan_treatment``).
+    """
+    if model == 'genome':
+        return registry.genome_kmers(model_dir).ab_list
+    if model == 'forecast':
+        return sorted(registry.lgbm(model_dir).vocabulary.get('antibiotics') or [])
+    if model == 'trend':
+        return sorted({ab for drugs in trend_series(model_dir).values() for ab in drugs})
+    if model == 'rl':
+        agent = registry.rl(model_dir)
+        return sorted(agent.drugs) if agent else []
+    raise ValueError("model must be 'genome', 'forecast', 'trend' or 'rl'")
+
+
+def status(model_dir=None):
+    """Report which models loaded, keyed by model name, with their measured
+    test metrics where they have them."""
+    agent = registry.rl(model_dir)
+    return {
+        'version': __version__,
+        'lgbm_forecasting': registry.lgbm(model_dir).status,
+        'genome_resistance': registry.genome(model_dir).status,
+        'trend_forecast': registry.trend(model_dir).status,
+        'rl_agent': {'trained': agent is not None,
+                     'model_type': agent.label if agent else None,
+                     'drugs': agent.drugs if agent else [],
+                     'description': 'Plans which drug to give each week; trained on the simulation'},
+        'mutation_timeline': registry.timeline(model_dir).status,
+    }
+
+
+def __getattr__(name):
+    # Re-export the classes without importing numpy/pandas/lightgbm at import time.
+    if name == 'LGBMResistancePredictor':
+        from .lgbm import LGBMResistancePredictor
+        return LGBMResistancePredictor
+    if name == 'GenomeModelPredictor':
+        from .genome import GenomeModelPredictor
+        return GenomeModelPredictor
+    if name == 'TimelineMLForecaster':
+        from .timeline_ml import TimelineMLForecaster
+        return TimelineMLForecaster
+    if name == 'RLAgent':
+        from .rl_agent import RLAgent
+        return RLAgent
+    if name == 'MutationTimelinePredictor':
+        from .timeline import MutationTimelinePredictor
+        return MutationTimelinePredictor
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

amrpredict/_paths.py

@@ -0,0 +1,11 @@
+"""Locating the model artifacts bundled inside the installed package."""
+from importlib.resources import files
+
+
+def default_model_dir():
+    """Absolute path to the models shipped with this package.
+
+    Used whenever a predictor is constructed without an explicit ``model_dir``,
+    so that ``LGBMForecaster()`` works straight after ``pip install``.
+    """
+    return str(files(__package__) / 'models')

amrpredict/amr_constants.py

@@ -0,0 +1,202 @@
+"""
+Antibiotic names and label maps: the library's copy of backend/amr_constants.py.
+
+The package installs without backend/, so it carries this copy. Do not edit it
+here: edit backend/amr_constants.py, then copy the data over (everything above
+FRONTEND_COPY); backend/tests/test_amr_constants.py fails while they differ.
+"""
+
+# Drug class per canonical antibiotic name. Anything missing is 'other'.
+DRUG_CLASS_MAP = {
+    'ampicillin': 'beta_lactam', 'amoxicillin': 'beta_lactam',
+    'amoxicillin/clavulanic acid': 'beta_lactam', 'piperacillin': 'beta_lactam',
+    'piperacillin/tazobactam': 'beta_lactam', 'oxacillin': 'beta_lactam',
+    'ampicillin/sulbactam': 'beta_lactam', 'penicillin': 'beta_lactam',
+    'methicillin': 'beta_lactam', 'carbenicillin': 'beta_lactam',
+    'cefazolin': 'beta_lactam', 'cefoxitin': 'beta_lactam', 'cefotaxime': 'beta_lactam',
+    'ceftazidime': 'beta_lactam', 'ceftriaxone': 'beta_lactam', 'cefepime': 'beta_lactam',
+    'cefuroxime': 'beta_lactam', 'cephalothin': 'beta_lactam', 'cefalothin': 'beta_lactam',
+    'ceftazidime/avibactam': 'beta_lactam', 'ceftolozane/tazobactam': 'beta_lactam',
+    'ceftiofur': 'beta_lactam', 'cefpodoxime': 'beta_lactam',
+    'imipenem': 'carbapenem', 'meropenem': 'carbapenem', 'ertapenem': 'carbapenem',
+    'doripenem': 'carbapenem', 'aztreonam': 'monobactam',
+    'ciprofloxacin': 'fluoroquinolone', 'levofloxacin': 'fluoroquinolone',
+    'norfloxacin': 'fluoroquinolone', 'nalidixic acid': 'fluoroquinolone',
+    'ofloxacin': 'fluoroquinolone', 'moxifloxacin': 'fluoroquinolone',
+    'pefloxacin': 'fluoroquinolone', 'trovafloxacin': 'fluoroquinolone',
+    'gentamicin': 'aminoglycoside', 'tobramycin': 'aminoglycoside',
+    'amikacin': 'aminoglycoside', 'streptomycin': 'aminoglycoside',
+    'neomycin': 'aminoglycoside', 'kanamycin': 'aminoglycoside',
+    'spectinomycin': 'aminoglycoside', 'capreomycin': 'aminoglycoside',
+    'tetracycline': 'tetracycline', 'doxycycline': 'tetracycline',
+    'minocycline': 'tetracycline', 'tigecycline': 'tetracycline',
+    'sulfamethoxazole': 'sulfonamide', 'trimethoprim': 'sulfonamide',
+    'trimethoprim/sulfamethoxazole': 'sulfonamide', 'sulfisoxazole': 'sulfonamide',
+    'chloramphenicol': 'phenicol',
+    'azithromycin': 'macrolide', 'erythromycin': 'macrolide',
+    'telithromycin': 'macrolide', 'clarithromycin': 'macrolide',
+    'colistin': 'polymyxin', 'polymyxin b': 'polymyxin',
+    'vancomycin': 'glycopeptide', 'teicoplanin': 'glycopeptide',
+    'clindamycin': 'lincosamide', 'nitrofurantoin': 'nitrofuran',
+    'rifampicin': 'rifamycin', 'rifampin': 'rifamycin', 'rifabutin': 'rifamycin',
+    'linezolid': 'oxazolidinone', 'daptomycin': 'lipopeptide',
+    'fusidic acid': 'fusidane', 'quinupristin/dalfopristin': 'streptogramin',
+    'pristinamycin': 'streptogramin', 'cephalexin': 'beta_lactam',
+    'cefpodoxime/clavulanic acid': 'beta_lactam', 'ticarcillin/clavulanic acid': 'beta_lactam',
+    'isoniazid': 'antitubercular', 'ethambutol': 'antitubercular',
+    'pyrazinamide': 'antitubercular', 'ethionamide': 'antitubercular',
+    'prothionamide': 'antitubercular', 'cycloserine': 'antitubercular',
+    'para-aminosalicylic acid': 'antitubercular', 'clofazimine': 'antitubercular',
+    'delamanid': 'antitubercular', 'bedaquiline': 'antitubercular',
+    # v4: drugs in the export (or the UI list) that fell into 'other'.
+    # Only azidothymidine (an antiviral) is left there on purpose.
+    'cefixime': 'beta_lactam', 'cefpirome': 'beta_lactam', 'cefoperazone': 'beta_lactam',
+    'cefotetan': 'beta_lactam', 'cefiderocol': 'beta_lactam', 'ceftaroline': 'beta_lactam',
+    'cefovecin': 'beta_lactam', 'cefamandole': 'beta_lactam', 'cefaclor': 'beta_lactam',
+    'cefmetazole': 'beta_lactam', 'ceftizoxime': 'beta_lactam', 'cefdinir': 'beta_lactam',
+    'cefozopran': 'beta_lactam', 'ceftobiprole': 'beta_lactam', 'ceftibuten': 'beta_lactam',
+    'ceftriaxone/cefpodoxime': 'beta_lactam', 'cefotaxime/clavulanic acid': 'beta_lactam',
+    'ceftazidime/clavulanic acid': 'beta_lactam', 'cefoperazone/sulbactam': 'beta_lactam',
+    'cefepime/taniborbactam': 'beta_lactam', 'temocillin': 'beta_lactam',
+    'ticarcillin': 'beta_lactam', 'mecillinam': 'beta_lactam', 'sulbactam': 'beta_lactam',
+    'imipenem/relebactam': 'carbapenem',
+    'netilmicin': 'aminoglycoside', 'plazomicin': 'aminoglycoside', 'apramycin': 'aminoglycoside',
+    'gatifloxacin': 'fluoroquinolone', 'enrofloxacin': 'fluoroquinolone',
+    'danofloxacin': 'fluoroquinolone', 'sparfloxacin': 'fluoroquinolone',
+    'pradofloxacin': 'fluoroquinolone', 'delafloxacin': 'fluoroquinolone',
+    'chlortetracycline': 'tetracycline', 'oxytetracycline': 'tetracycline',
+    'eravacycline': 'tetracycline', 'omadacycline': 'tetracycline',
+    'sulfathiazole': 'sulfonamide', 'sulfamethazine': 'sulfonamide',
+    'trimethoprim/sulfobactam': 'sulfonamide',
+    'tylosin': 'macrolide', 'tulathromycin': 'macrolide', 'spiramycin': 'macrolide',
+    'florfenicol': 'phenicol', 'lincomycin': 'lincosamide', 'virginiamycin': 'streptogramin',
+    'furazolidone': 'nitrofuran', 'metronidazole': 'nitroimidazole',
+    'fosfomycin': 'phosphonic_acid', 'mupirocin': 'pseudomonic_acid',
+    'zoliflodacin': 'spiropyrimidinetrione', 'avilamycin': 'orthosomycin',
+    'nicotinamide': 'antitubercular',   # all 229 rows are Mycobacterium
+}
+
+# Spelling variants, typos and non-drugs found in the raw Antibiotic column.
+# Left-hand side is what appears in the export; right-hand side is canonical.
+# None means "drop these rows", a drug class is not a drug.
+ANTIBIOTIC_ALIASES = {
+    'ampicillin-sulbactam': 'ampicillin/sulbactam',
+    'ampicillin_clavulanic_acid': 'amoxicillin/clavulanic acid',
+    'amoxicillin-clavulanic acid': 'amoxicillin/clavulanic acid',
+    'piperacillin-tazobactam': 'piperacillin/tazobactam',
+    'trimethoprim-sulfamethoxazole': 'trimethoprim/sulfamethoxazole',
+    'sulfamethoxazole/trimethoprim': 'trimethoprim/sulfamethoxazole',
+    'co_trimoxazole': 'trimethoprim/sulfamethoxazole',
+    'co-trimoxazole': 'trimethoprim/sulfamethoxazole',
+    'geamycin': 'gentamicin',
+    'trimotheprim': 'trimethoprim',
+    'cefalothin': 'cephalothin',
+    'rifampin': 'rifampicin',
+    # Separator variants of combination drugs
+    'tazobactam_piperacillin': 'piperacillin/tazobactam',
+    'ceftazidime_avibactam': 'ceftazidime/avibactam',
+    'ceftolozane_tazobactam': 'ceftolozane/tazobactam',
+    'ticarcillin_clavulanate': 'ticarcillin/clavulanic acid',
+    'cefpodoxime_clavulanic_acid': 'cefpodoxime/clavulanic acid',
+    'trimethoprim_sulfobactam': 'trimethoprim/sulfobactam',
+    'para_aminosalicylic_acid': 'para-aminosalicylic acid',
+    'cefepime_taniborbactam': 'cefepime/taniborbactam',
+    'amoxicillin_clavulanat': 'amoxicillin/clavulanic acid',
+    'polymyxin_b': 'polymyxin b',
+    'ceftazidime-avibactam': 'ceftazidime/avibactam',
+    'ceftolozane-tazobactam': 'ceftolozane/tazobactam',
+    'imipenem-relebactam': 'imipenem/relebactam',
+    # Typos and a broken character encoding
+    'amipicillin_sulbactam': 'ampicillin/sulbactam',
+    'tgecycline': 'tigecycline',
+    'cefuroxim\u00e2': 'cefuroxime',
+    'pristimycin': 'pristinamycin',
+    'cefotaxime/clavulanic acid\u00e2': 'cefotaxime/clavulanic acid',
+    # Other-language spellings, found by Suleman in BVBRC_genome_amr.csv
+    'tigecyklin': 'tigecycline',
+    'tetracyklin': 'tetracycline',
+    'cefpirom': 'cefpirome',
+    # Shigella panel; never on the same genome as nitrofurantoin
+    'strofurantoin': 'nitrofurantoin',
+    # Alternative names for the same drug
+    'cefuroxime_sodium': 'cefuroxime',
+    'cefalotin': 'cephalothin',
+    'cefalexin': 'cephalexin',
+    'synercid': 'quinupristin/dalfopristin',
+    'phosphomycin': 'fosfomycin',
+    'benzylpenicillin': 'penicillin',
+    # Not a single drug: drug classes, a phenotype, and a lost drug name
+    # ('instrument' is 593 C. difficile lab rows with the drug name missing)
+    'carbapenem': None,
+    'beta-lactam': None,
+    'cephalosporin': None,
+    'fluoroquinolones': None,
+    'aminogycosides': None,
+    'macrolides': None,
+    'sulfonamides': None,
+    'sulfa': None,
+    'extended spectrum beta lactamase': None,
+    'instrument': None,
+}
+
+# Resistant Phenotype -> label. 0 = susceptible, 1 = resistant.
+PHENOTYPE_MAP = {
+    'Susceptible': 0,
+    'Susceptible-dose dependent': 0,
+    'Resistant': 1,
+    'Intermediate': 1,
+    'Nonsusceptible': 1,
+    'Reduced Susceptibility': 1,
+}
+
+# Phenotypes each trainer keeps; the rest of its rows are dropped. They
+# differ on purpose, to keep the served models' training data as it was.
+LGBM_TRAIN_PHENOTYPES = ('Susceptible', 'Resistant', 'Intermediate', 'Nonsusceptible')
+KMER_TRAIN_PHENOTYPES = ('Susceptible', 'Resistant', 'Intermediate')
+
+# MIC comparators a user may type that the forecaster has no category for,
+# mapped onto the nearest sign it learned.
+MIC_SIGN_ALIASES = {'>=': '>', '\u2265': '>', '\u2264': '<=', '=<': '<=', '=>': '>'}
+
+# Antibiotics offered in the dropdowns, grouped by class. The pages show only
+# the names the loaded model knows; this list is the fallback when the backend
+# is down, and what GET /api/antibiotics/ returns.
+UI_ANTIBIOTICS = [
+    'penicillin', 'ampicillin', 'ampicillin/sulbactam', 'amoxicillin',
+    'amoxicillin/clavulanic acid', 'piperacillin', 'piperacillin/tazobactam',
+    'oxacillin', 'temocillin', 'carbenicillin', 'ticarcillin/clavulanic acid',
+    'cefazolin', 'cefoxitin', 'cefotetan', 'cefmetazole', 'cefotaxime',
+    'cefotaxime/clavulanic acid', 'ceftazidime', 'ceftazidime/avibactam',
+    'ceftazidime/clavulanic acid', 'ceftolozane/tazobactam', 'ceftriaxone',
+    'cefepime', 'cefepime/taniborbactam', 'cefuroxime', 'cephalothin',
+    'cefixime', 'cefpodoxime', 'cefpodoxime/clavulanic acid', 'ceftibuten',
+    'cefoperazone/sulbactam', 'ceftiofur', 'cefpirome', 'cefozopran',
+    'ceftaroline', 'ceftobiprole',
+    'imipenem', 'imipenem/relebactam', 'meropenem', 'ertapenem', 'doripenem',
+    'aztreonam', 'ciprofloxacin', 'levofloxacin', 'norfloxacin',
+    'nalidixic acid', 'ofloxacin', 'moxifloxacin', 'pefloxacin',
+    'delafloxacin',
+    'gentamicin', 'tobramycin', 'amikacin', 'streptomycin', 'neomycin',
+    'kanamycin', 'spectinomycin', 'apramycin',
+    'tetracycline', 'oxytetracycline', 'doxycycline', 'minocycline',
+    'tigecycline',
+    'sulfamethoxazole', 'sulfisoxazole', 'trimethoprim',
+    'trimethoprim/sulfamethoxazole',
+    'chloramphenicol', 'florfenicol',
+    'azithromycin', 'erythromycin', 'clarithromycin', 'telithromycin',
+    'colistin', 'polymyxin b', 'vancomycin', 'teicoplanin',
+    'clindamycin', 'lincomycin', 'nitrofurantoin', 'fosfomycin', 'rifampicin',
+]
+
+
+def normalize_antibiotic(name):
+    """Lower-case, strip, and map a spelling variant to its canonical name.
+
+    Names the alias table drops from training (drug classes such as
+    'carbapenem') are returned unchanged; a model then reports them as
+    unrecognised instead of silently answering for a different drug.
+    """
+    if name in (None, ''):
+        return name
+    ab = ' '.join(str(name).split()).lower()
+    return ANTIBIOTIC_ALIASES.get(ab) or ab

amrpredict/cli.py

@@ -0,0 +1,125 @@
+"""Command-line interface: ``amrpredict <command> [options]``."""
+import argparse
+import json
+import sys
+
+
+def _read_fasta(path):
+    if path == '-':
+        return sys.stdin.read()
+    with open(path, 'r', encoding='utf-8', errors='replace') as fh:
+        return fh.read()
+
+
+def _emit(payload, pretty):
+    print(json.dumps(payload, indent=2 if pretty else None, default=str))
+
+
+def build_parser():
+    from . import __version__
+
+    parser = argparse.ArgumentParser(
+        prog='amrpredict',
+        description='Predict antimicrobial resistance from genomic data.',
+    )
+    parser.add_argument('--version', action='version', version=f'amrpredict {__version__}')
+    parser.add_argument('--compact', action='store_true', help='emit single-line JSON')
+    parser.add_argument('--model-dir', default=None, help='use artifacts from this directory')
+    sub = parser.add_subparsers(dest='command', required=True)
+
+    p = sub.add_parser('forecast', help='predict from metadata + MIC value')
+    p.add_argument('antibiotic')
+    p.add_argument('--taxon-id', type=int, default=None, help='NCBI taxonomy ID, e.g. 562')
+    p.add_argument('--mic-value', type=float, default=None, help='MIC in mg/L')
+    p.add_argument('--mic-sign', default=None, help="MIC comparator, e.g. '=' or '>'")
+    p.add_argument('--genus', default='unknown')
+    p.add_argument('--species', default='unknown')
+    p.add_argument('--threshold', type=float, default=None,
+                   help="default: the bundled model's own threshold")
+
+    p = sub.add_parser('predict', help='predict from a complete genome FASTA')
+    p.add_argument('fasta', help="path to a FASTA file, or '-' for stdin")
+    p.add_argument('antibiotic')
+    p.add_argument('--threshold', type=float, default=None,
+                   help="default: the model's own threshold")
+    p.add_argument('--genus', default=None)
+    p.add_argument('--species', default=None)
+    genes = p.add_mutually_exclusive_group()
+    genes.add_argument('--kmers', action='store_true', help='always use the k-mer model')
+    genes.add_argument('--genes', action='store_true',
+                       help='require the gene model (needs AMRFinderPlus)')
+
+    p = sub.add_parser('species', help='identify the species of a complete genome')
+    p.add_argument('fasta', help="path to a FASTA file, or '-' for stdin")
+
+    p = sub.add_parser('trend', help='yearly resistance of a genus and drug, and its forecast')
+    p.add_argument('genus')
+    p.add_argument('antibiotic')
+    p.add_argument('--years', type=int, default=5, help='years ahead, 1 to 10')
+
+    p = sub.add_parser('plan', help='drug per week: the RL agent against fixed schedules')
+    p.add_argument('antibiotic', help='ciprofloxacin, gentamicin or imipenem')
+    p.add_argument('--weeks', type=int, default=52, help='1 to 104')
+    p.add_argument('--seed', type=int, default=42)
+
+    p = sub.add_parser('timeline', help='simulate resistance evolution (not a trained model)')
+    p.add_argument('fasta', help="path to a FASTA file, or '-' for stdin")
+    p.add_argument('antibiotic')
+    p.add_argument('--weeks', type=int, default=8)
+
+    p = sub.add_parser('antibiotics', help='list the antibiotics a model knows')
+    p.add_argument('--model', default='genome', choices=['genome', 'forecast', 'trend', 'rl'])
+    sub.add_parser('status', help='report which models loaded')
+    return parser
+
+
+def main(argv=None):
+    import amrpredict
+
+    args = build_parser().parse_args(argv)
+    pretty = not args.compact
+    md = args.model_dir
+
+    try:
+        if args.command == 'forecast':
+            result = amrpredict.forecast(
+                args.antibiotic, taxon_id=args.taxon_id, mic_value=args.mic_value,
+                mic_sign=args.mic_sign, genus=args.genus, species=args.species,
+                threshold=args.threshold, model_dir=md,
+            )
+        elif args.command == 'predict':
+            result = amrpredict.predict_fasta(
+                _read_fasta(args.fasta), args.antibiotic, threshold=args.threshold,
+                genus=args.genus, species=args.species,
+                use_genes=False if args.kmers else (True if args.genes else None), model_dir=md,
+            )
+        elif args.command == 'species':
+            result = amrpredict.identify_species(_read_fasta(args.fasta), model_dir=md)
+        elif args.command == 'trend':
+            result = amrpredict.forecast_trend(args.genus, args.antibiotic,
+                                               years_ahead=args.years, model_dir=md)
+        elif args.command == 'plan':
+            result = amrpredict.plan_treatment(args.antibiotic, n_weeks=args.weeks,
+                                               seed=args.seed, model_dir=md)
+        elif args.command == 'timeline':
+            result = amrpredict.simulate_timeline(
+                _read_fasta(args.fasta), args.antibiotic,
+                n_weeks=args.weeks, model_dir=md,
+            )
+        elif args.command == 'antibiotics':
+            result = amrpredict.antibiotics(args.model, model_dir=md)
+        else:
+            result = amrpredict.status(model_dir=md)
+    except FileNotFoundError as exc:
+        print(f'amrpredict: {exc}', file=sys.stderr)
+        return 2
+    except ValueError as exc:
+        print(f'amrpredict: {exc}', file=sys.stderr)
+        return 1
+
+    _emit(result, pretty)
+    return 0
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())