Initial commit
This commit is contained in:
Zhongwei Li
447
skills/clinical-decision-support/scripts/build_decision_tree.py
Executable file
447
skills/clinical-decision-support/scripts/build_decision_tree.py
Executable file
@@ -0,0 +1,447 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Build Clinical Decision Tree Flowcharts in TikZ Format
|
||||
|
||||
Generates LaTeX/TikZ code for clinical decision algorithms from
|
||||
simple text or YAML descriptions.
|
||||
|
||||
Dependencies: pyyaml (optional, for YAML input)
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
|
||||
class DecisionNode:
|
||||
"""Represents a decision point in the clinical algorithm."""
|
||||
|
||||
def __init__(self, question, yes_path=None, no_path=None, node_id=None):
|
||||
self.question = question
|
||||
self.yes_path = yes_path
|
||||
self.no_path = no_path
|
||||
self.node_id = node_id or self._generate_id(question)
|
||||
|
||||
def _generate_id(self, text):
|
||||
"""Generate clean node ID from text."""
|
||||
return ''.join(c for c in text if c.isalnum())[:15].lower()
|
||||
|
||||
|
||||
class ActionNode:
|
||||
"""Represents an action/outcome in the clinical algorithm."""
|
||||
|
||||
def __init__(self, action, urgency='routine', node_id=None):
|
||||
self.action = action
|
||||
self.urgency = urgency # 'urgent', 'semiurgent', 'routine'
|
||||
self.node_id = node_id or self._generate_id(action)
|
||||
|
||||
def _generate_id(self, text):
|
||||
return ''.join(c for c in text if c.isalnum())[:15].lower()
|
||||
|
||||
|
||||
def generate_tikz_header():
|
||||
"""Generate TikZ preamble with style definitions."""
|
||||
|
||||
tikz = """\\documentclass[10pt]{article}
|
||||
\\usepackage[margin=0.5in, landscape]{geometry}
|
||||
\\usepackage{tikz}
|
||||
\\usetikzlibrary{shapes,arrows,positioning}
|
||||
\\usepackage{xcolor}
|
||||
|
||||
% Color definitions
|
||||
\\definecolor{urgentred}{RGB}{220,20,60}
|
||||
\\definecolor{actiongreen}{RGB}{0,153,76}
|
||||
\\definecolor{decisionyellow}{RGB}{255,193,7}
|
||||
\\definecolor{routineblue}{RGB}{100,181,246}
|
||||
\\definecolor{headerblue}{RGB}{0,102,204}
|
||||
|
||||
% TikZ styles
|
||||
\\tikzstyle{startstop} = [rectangle, rounded corners=8pt, minimum width=3cm, minimum height=1cm,
|
||||
text centered, draw=black, fill=headerblue!20, font=\\small\\bfseries]
|
||||
\\tikzstyle{decision} = [diamond, minimum width=3cm, minimum height=1.2cm, text centered,
|
||||
draw=black, fill=decisionyellow!40, font=\\small, aspect=2, inner sep=0pt,
|
||||
text width=3.5cm]
|
||||
\\tikzstyle{process} = [rectangle, rounded corners=4pt, minimum width=3.5cm, minimum height=0.9cm,
|
||||
text centered, draw=black, fill=actiongreen!20, font=\\small]
|
||||
\\tikzstyle{urgent} = [rectangle, rounded corners=4pt, minimum width=3.5cm, minimum height=0.9cm,
|
||||
text centered, draw=urgentred, line width=1.5pt, fill=urgentred!15,
|
||||
font=\\small\\bfseries]
|
||||
\\tikzstyle{routine} = [rectangle, rounded corners=4pt, minimum width=3.5cm, minimum height=0.9cm,
|
||||
text centered, draw=black, fill=routineblue!20, font=\\small]
|
||||
\\tikzstyle{arrow} = [thick,->,>=stealth]
|
||||
\\tikzstyle{urgentarrow} = [ultra thick,->,>=stealth,color=urgentred]
|
||||
|
||||
\\begin{document}
|
||||
|
||||
\\begin{center}
|
||||
{\\Large\\bfseries Clinical Decision Algorithm}\\\\[10pt]
|
||||
{\\large [TITLE TO BE SPECIFIED]}
|
||||
\\end{center}
|
||||
|
||||
\\vspace{10pt}
|
||||
|
||||
\\begin{tikzpicture}[node distance=2.2cm and 3.5cm, auto]
|
||||
|
||||
"""
|
||||
|
||||
return tikz
|
||||
|
||||
|
||||
def generate_tikz_footer():
|
||||
"""Generate TikZ closing code."""
|
||||
|
||||
tikz = """
|
||||
\\end{tikzpicture}
|
||||
|
||||
\\end{document}
|
||||
"""
|
||||
|
||||
return tikz
|
||||
|
||||
|
||||
def simple_algorithm_to_tikz(algorithm_text, output_file='algorithm.tex'):
|
||||
"""
|
||||
Convert simple text-based algorithm to TikZ flowchart.
|
||||
|
||||
Input format (simple question-action pairs):
|
||||
START: Chief complaint
|
||||
Q1: High-risk criteria present? -> YES: Immediate action (URGENT) | NO: Continue
|
||||
Q2: Risk score >= 3? -> YES: Admit ICU | NO: Outpatient management (ROUTINE)
|
||||
END: Final outcome
|
||||
|
||||
Parameters:
|
||||
algorithm_text: Multi-line string with algorithm
|
||||
output_file: Path to save .tex file
|
||||
"""
|
||||
|
||||
tikz_code = generate_tikz_header()
|
||||
|
||||
# Parse algorithm text
|
||||
lines = [line.strip() for line in algorithm_text.strip().split('\n') if line.strip()]
|
||||
|
||||
node_defs = []
|
||||
arrow_defs = []
|
||||
|
||||
previous_node = None
|
||||
node_counter = 0
|
||||
|
||||
for line in lines:
|
||||
if line.startswith('START:'):
|
||||
# Start node
|
||||
text = line.replace('START:', '').strip()
|
||||
node_id = 'start'
|
||||
node_defs.append(f"\\node [startstop] ({node_id}) {{{text}}};")
|
||||
previous_node = node_id
|
||||
node_counter += 1
|
||||
|
||||
elif line.startswith('END:'):
|
||||
# End node
|
||||
text = line.replace('END:', '').strip()
|
||||
node_id = 'end'
|
||||
|
||||
# Position relative to previous
|
||||
if previous_node:
|
||||
node_defs.append(f"\\node [startstop, below=of {previous_node}] ({node_id}) {{{text}}};")
|
||||
arrow_defs.append(f"\\draw [arrow] ({previous_node}) -- ({node_id});")
|
||||
|
||||
elif line.startswith('Q'):
|
||||
# Decision node
|
||||
parts = line.split(':', 1)
|
||||
if len(parts) < 2:
|
||||
continue
|
||||
|
||||
question_part = parts[1].split('->')[0].strip()
|
||||
node_id = f'q{node_counter}'
|
||||
|
||||
# Add decision node
|
||||
if previous_node:
|
||||
node_defs.append(f"\\node [decision, below=of {previous_node}] ({node_id}) {{{question_part}}};")
|
||||
arrow_defs.append(f"\\draw [arrow] ({previous_node}) -- ({node_id});")
|
||||
else:
|
||||
node_defs.append(f"\\node [decision] ({node_id}) {{{question_part}}};")
|
||||
|
||||
# Parse YES and NO branches
|
||||
if '->' in line:
|
||||
branches = line.split('->')[1].split('|')
|
||||
|
||||
for branch in branches:
|
||||
branch = branch.strip()
|
||||
|
||||
if branch.startswith('YES:'):
|
||||
yes_action = branch.replace('YES:', '').strip()
|
||||
yes_id = f'yes{node_counter}'
|
||||
|
||||
# Check urgency
|
||||
if '(URGENT)' in yes_action:
|
||||
style = 'urgent'
|
||||
yes_action = yes_action.replace('(URGENT)', '').strip()
|
||||
arrow_style = 'urgentarrow'
|
||||
elif '(ROUTINE)' in yes_action:
|
||||
style = 'routine'
|
||||
yes_action = yes_action.replace('(ROUTINE)', '').strip()
|
||||
arrow_style = 'arrow'
|
||||
else:
|
||||
style = 'process'
|
||||
arrow_style = 'arrow'
|
||||
|
||||
node_defs.append(f"\\node [{style}, left=of {node_id}] ({yes_id}) {{{yes_action}}};")
|
||||
arrow_defs.append(f"\\draw [{arrow_style}] ({node_id}) -- node[above] {{Yes}} ({yes_id});")
|
||||
|
||||
elif branch.startswith('NO:'):
|
||||
no_action = branch.replace('NO:', '').strip()
|
||||
no_id = f'no{node_counter}'
|
||||
|
||||
# Check urgency
|
||||
if '(URGENT)' in no_action:
|
||||
style = 'urgent'
|
||||
no_action = no_action.replace('(URGENT)', '').strip()
|
||||
arrow_style = 'urgentarrow'
|
||||
elif '(ROUTINE)' in no_action:
|
||||
style = 'routine'
|
||||
no_action = no_action.replace('(ROUTINE)', '').strip()
|
||||
arrow_style = 'arrow'
|
||||
else:
|
||||
style = 'process'
|
||||
arrow_style = 'arrow'
|
||||
|
||||
node_defs.append(f"\\node [{style}, right=of {node_id}] ({no_id}) {{{no_action}}};")
|
||||
arrow_defs.append(f"\\draw [{arrow_style}] ({node_id}) -- node[above] {{No}} ({no_id});")
|
||||
|
||||
previous_node = node_id
|
||||
node_counter += 1
|
||||
|
||||
# Add all nodes and arrows to TikZ
|
||||
tikz_code += '\n'.join(node_defs) + '\n\n'
|
||||
tikz_code += '% Arrows\n'
|
||||
tikz_code += '\n'.join(arrow_defs) + '\n'
|
||||
|
||||
tikz_code += generate_tikz_footer()
|
||||
|
||||
# Save to file
|
||||
with open(output_file, 'w') as f:
|
||||
f.write(tikz_code)
|
||||
|
||||
print(f"TikZ flowchart saved to: {output_file}")
|
||||
print(f"Compile with: pdflatex {output_file}")
|
||||
|
||||
return tikz_code
|
||||
|
||||
|
||||
def json_to_tikz(json_file, output_file='algorithm.tex'):
|
||||
"""
|
||||
Convert JSON decision tree specification to TikZ flowchart.
|
||||
|
||||
JSON format:
|
||||
{
|
||||
"title": "Algorithm Title",
|
||||
"nodes": {
|
||||
"start": {"type": "start", "text": "Patient presentation"},
|
||||
"q1": {"type": "decision", "text": "Criteria met?", "yes": "action1", "no": "q2"},
|
||||
"action1": {"type": "action", "text": "Immediate intervention", "urgency": "urgent"},
|
||||
"q2": {"type": "decision", "text": "Score >= 3?", "yes": "action2", "no": "action3"},
|
||||
"action2": {"type": "action", "text": "Admit ICU"},
|
||||
"action3": {"type": "action", "text": "Outpatient", "urgency": "routine"}
|
||||
},
|
||||
"start_node": "start"
|
||||
}
|
||||
"""
|
||||
|
||||
with open(json_file, 'r') as f:
|
||||
spec = json.load(f)
|
||||
|
||||
tikz_code = generate_tikz_header()
|
||||
|
||||
# Replace title
|
||||
title = spec.get('title', 'Clinical Decision Algorithm')
|
||||
tikz_code = tikz_code.replace('[TITLE TO BE SPECIFIED]', title)
|
||||
|
||||
nodes = spec['nodes']
|
||||
start_node = spec.get('start_node', 'start')
|
||||
|
||||
# Generate nodes (simplified layout - vertical)
|
||||
node_defs = []
|
||||
arrow_defs = []
|
||||
|
||||
# Track positioning
|
||||
previous_node = None
|
||||
level = 0
|
||||
|
||||
def add_node(node_id, position_rel=None):
|
||||
"""Recursively add nodes."""
|
||||
|
||||
if node_id not in nodes:
|
||||
return
|
||||
|
||||
node = nodes[node_id]
|
||||
node_type = node['type']
|
||||
text = node['text']
|
||||
|
||||
# Determine TikZ style
|
||||
if node_type == 'start' or node_type == 'end':
|
||||
style = 'startstop'
|
||||
elif node_type == 'decision':
|
||||
style = 'decision'
|
||||
elif node_type == 'action':
|
||||
urgency = node.get('urgency', 'normal')
|
||||
if urgency == 'urgent':
|
||||
style = 'urgent'
|
||||
elif urgency == 'routine':
|
||||
style = 'routine'
|
||||
else:
|
||||
style = 'process'
|
||||
else:
|
||||
style = 'process'
|
||||
|
||||
# Position node
|
||||
if position_rel:
|
||||
node_def = f"\\node [{style}, {position_rel}] ({node_id}) {{{text}}};"
|
||||
else:
|
||||
node_def = f"\\node [{style}] ({node_id}) {{{text}}};"
|
||||
|
||||
node_defs.append(node_def)
|
||||
|
||||
# Add arrows for decision nodes
|
||||
if node_type == 'decision':
|
||||
yes_target = node.get('yes')
|
||||
no_target = node.get('no')
|
||||
|
||||
if yes_target:
|
||||
# Determine arrow style based on target urgency
|
||||
target_node = nodes.get(yes_target, {})
|
||||
arrow_style = 'urgentarrow' if target_node.get('urgency') == 'urgent' else 'arrow'
|
||||
arrow_defs.append(f"\\draw [{arrow_style}] ({node_id}) -| node[near start, above] {{Yes}} ({yes_target});")
|
||||
|
||||
if no_target:
|
||||
target_node = nodes.get(no_target, {})
|
||||
arrow_style = 'urgentarrow' if target_node.get('urgency') == 'urgent' else 'arrow'
|
||||
arrow_defs.append(f"\\draw [{arrow_style}] ({node_id}) -| node[near start, above] {{No}} ({no_target});")
|
||||
|
||||
# Simple layout - just list nodes (manual positioning in JSON works better for complex trees)
|
||||
for node_id in nodes.keys():
|
||||
add_node(node_id)
|
||||
|
||||
tikz_code += '\n'.join(node_defs) + '\n\n'
|
||||
tikz_code += '% Arrows\n'
|
||||
tikz_code += '\n'.join(arrow_defs) + '\n'
|
||||
|
||||
tikz_code += generate_tikz_footer()
|
||||
|
||||
# Save
|
||||
with open(output_file, 'w') as f:
|
||||
f.write(tikz_code)
|
||||
|
||||
print(f"TikZ flowchart saved to: {output_file}")
|
||||
return tikz_code
|
||||
|
||||
|
||||
def create_example_json():
|
||||
"""Create example JSON specification for testing."""
|
||||
|
||||
example = {
|
||||
"title": "Acute Chest Pain Management Algorithm",
|
||||
"nodes": {
|
||||
"start": {
|
||||
"type": "start",
|
||||
"text": "Patient with\\nchest pain"
|
||||
},
|
||||
"q1": {
|
||||
"type": "decision",
|
||||
"text": "STEMI\\ncriteria?",
|
||||
"yes": "stemi_action",
|
||||
"no": "q2"
|
||||
},
|
||||
"stemi_action": {
|
||||
"type": "action",
|
||||
"text": "Activate cath lab\\nAspirin, heparin\\nPrimary PCI",
|
||||
"urgency": "urgent"
|
||||
},
|
||||
"q2": {
|
||||
"type": "decision",
|
||||
"text": "High-risk\\nfeatures?",
|
||||
"yes": "admit",
|
||||
"no": "q3"
|
||||
},
|
||||
"admit": {
|
||||
"type": "action",
|
||||
"text": "Admit CCU\\nSerial troponins\\nEarly angiography"
|
||||
},
|
||||
"q3": {
|
||||
"type": "decision",
|
||||
"text": "TIMI\\nscore 0-1?",
|
||||
"yes": "lowrisk",
|
||||
"no": "moderate"
|
||||
},
|
||||
"lowrisk": {
|
||||
"type": "action",
|
||||
"text": "Observe 6-12h\\nStress test\\nOutpatient f/u",
|
||||
"urgency": "routine"
|
||||
},
|
||||
"moderate": {
|
||||
"type": "action",
|
||||
"text": "Admit telemetry\\nMedical management\\nRisk stratification"
|
||||
}
|
||||
},
|
||||
"start_node": "start"
|
||||
}
|
||||
|
||||
return example
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Build clinical decision tree flowcharts')
|
||||
parser.add_argument('-i', '--input', type=str, default=None,
|
||||
help='Input file (JSON format)')
|
||||
parser.add_argument('-o', '--output', type=str, default='clinical_algorithm.tex',
|
||||
help='Output .tex file')
|
||||
parser.add_argument('--example', action='store_true',
|
||||
help='Generate example algorithm')
|
||||
parser.add_argument('--text', type=str, default=None,
|
||||
help='Simple text algorithm (see format in docstring)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.example:
|
||||
print("Generating example algorithm...")
|
||||
example_spec = create_example_json()
|
||||
|
||||
# Save example JSON
|
||||
with open('example_algorithm.json', 'w') as f:
|
||||
json.dump(example_spec, f, indent=2)
|
||||
print("Example JSON saved to: example_algorithm.json")
|
||||
|
||||
# Generate TikZ from example
|
||||
json_to_tikz('example_algorithm.json', args.output)
|
||||
|
||||
elif args.text:
|
||||
print("Generating algorithm from text...")
|
||||
simple_algorithm_to_tikz(args.text, args.output)
|
||||
|
||||
elif args.input:
|
||||
print(f"Generating algorithm from {args.input}...")
|
||||
if args.input.endswith('.json'):
|
||||
json_to_tikz(args.input, args.output)
|
||||
else:
|
||||
with open(args.input, 'r') as f:
|
||||
text = f.read()
|
||||
simple_algorithm_to_tikz(text, args.output)
|
||||
|
||||
else:
|
||||
print("No input provided. Use --example to generate example, --text for simple text, or -i for JSON input.")
|
||||
print("\nSimple text format:")
|
||||
print("START: Patient presentation")
|
||||
print("Q1: Criteria met? -> YES: Action (URGENT) | NO: Continue")
|
||||
print("Q2: Score >= 3? -> YES: Admit | NO: Outpatient (ROUTINE)")
|
||||
print("END: Follow-up")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
|
||||
# Example usage:
|
||||
# python build_decision_tree.py --example
|
||||
# python build_decision_tree.py -i algorithm_spec.json -o my_algorithm.tex
|
||||
#
|
||||
# Then compile:
|
||||
# pdflatex clinical_algorithm.tex
|
||||
|
||||
Reference in New Issue
Block a user