alethiotx.artemis.mesh

 5def tree(s3_base = None, url_base = None, file_base = None) -> dict:
 6    """
 7    Retrieve MeSH tree structure from S3 cache or download from URL.
 8    This function attempts to load the MeSH (Medical Subject Headings) tree structure,
 9    first trying an S3 cache if provided, then falling back to downloading from a URL.
10    :param s3_base: Base path for S3 storage location (optional). If provided, ``file_base`` is required.
11    :type s3_base: str or None
12    :param url_base: Base URL for downloading MeSH tree file (optional). Requires ``file_base``.
13    :type url_base: str or None
14    :param file_base: Base filename (without extension) for the MeSH tree file (optional).
15                      Required when either ``s3_base`` or ``url_base`` is provided.
16    :type file_base: str or None
17    :return: Dictionary mapping MeSH headings to their tree numbers.
18             Returns empty dict if all retrieval methods fail.
19    :rtype: dict
20    :raises: Prints error messages to console but does not raise exceptions.
21    .. note::
22       - Priority order: S3 cache → URL download
23       - S3 files should be pickle format (.pkl extension)
24       - URL files should be MeSH ASCII format (.bin extension)
25       - MeSH ASCII format uses 'MH =' for headings and 'MN =' for tree numbers
26
27    **Example**
28
29    >>> mesh = tree(s3_base='s3://alethiotx-artemis/data/mesh/', file_base='d2025')
30    >>> mesh = tree(url_base='https://nlmpubs.nlm.nih.gov/projects/mesh/MESH_FILES/asciimesh/', file_base='d2025')
31    """
32    # Prioritize S3 if s3_base is provided
33    if s3_base is not None:
34        if file_base is None:
35            print("Error: file_base is required when s3_base is provided")
36            return {}
37        try:
38            print(f"Checking for cached MeSH tree at {s3_base}...")
39            mesh_tree = read_pickle(s3_base + file_base + ".pkl")
40            print(f"MeSH tree loaded from S3: {len(mesh_tree)} headings")
41            return mesh_tree
42        except Exception as e:
43            print(f"S3 cache not found or could not be loaded: {e}")
44            # Fall through to try URL download if available
45    
46    # If S3 failed or not provided, try URL download
47    if url_base is not None and file_base is not None:
48        # Continue with downloading from URL
49        mesh_url = url_base + file_base + ".bin"
50        print(f"Downloading MeSH tree from {mesh_url}...")
51    
52        try:
53            response = requests.get(mesh_url, timeout=60)
54            response.raise_for_status()
55            
56            # Parse the MeSH ASCII file
57            mesh_tree = {}
58            current_heading = None
59            current_trees = []
60            
61            for line in response.text.split('\n'):
62                line = line.strip()
63                
64                # MH = MeSH Heading
65                if line.startswith('MH = '):
66                    current_heading = line[5:].strip()
67                    current_trees = []
68                
69                # MN = Tree Number
70                elif line.startswith('MN = '):
71                    tree_num = line[5:].strip()
72                    current_trees.append(tree_num)
73                
74                # Entry = end of record, store the data
75                elif line.startswith('ENTRY') or line.startswith('*NEWRECORD'):
76                    if current_heading and current_trees:
77                        mesh_tree[current_heading] = current_trees
78                        current_heading = None
79                        current_trees = []
80            
81            print(f"MeSH tree loaded: {len(mesh_tree)} headings")
82            
83            return mesh_tree
84            
85        except Exception as e:
86            print(f"Warning: Could not download MeSH tree: {e}")
87            return {}
88    
89    # Neither S3 nor URL options are available or both failed
90    if s3_base is not None:
91        print("Error: S3 download failed and no valid URL options provided")
92    else:
93        print("Error: Either (url_base and file_base) or (s3_base and file_base) are required")
94    return {}

 96def descendants(heading: str, s3_base=None, file_base=None, url_base = None) -> list:
 97    """
 98    Get all descendant MeSH headings for a given heading based on tree hierarchy.
 99    This function retrieves a heading and all of its descendants in the MeSH tree
100    structure by comparing tree numbers. A heading is considered a descendant if
101    its tree number starts with a parent's tree number followed by a dot.
102    :param heading: The MeSH heading to find descendants for
103    :type heading: str
104    :param s3_base: S3 bucket base path for MeSH data, defaults to None
105    :type s3_base: str, optional
106    :param file_base: Local file system base path for MeSH data, defaults to None
107    :type file_base: str, optional
108    :param url_base: URL base path for MeSH data, defaults to None
109    :type url_base: str, optional
110    :return: List of MeSH headings including the parent and all descendants
111    :rtype: list
112    :raises: None - Returns list with single heading if no tree numbers found
113    .. note::
114        The function uses tree numbers to determine hierarchy. A tree number like
115        'A01.1.2' is a descendant of 'A01.1'.
116
117    **Example**
118
119    >>> descendants('Nervous System')
120    ['Nervous System', 'Central Nervous System', 'Brain', ...]
121    """
122    # Get the full tree info (headings with their tree numbers)
123    full_tree = tree(s3_base = s3_base, url_base = url_base, file_base=file_base)
124    
125    # Get tree numbers for the requested heading
126    parent_trees = full_tree.get(heading)
127    if not parent_trees:
128        print(f"No tree numbers found for '{heading}', using exact match only")
129        return [heading]
130    
131    print(f"Tree numbers for '{heading}': {parent_trees}")
132    
133    # Find all descendants recursively based on tree number hierarchy
134    descendants = set([heading])
135    
136    # Check all headings to see if any of their tree numbers are descendants
137    for heading, tree_nums in full_tree.items():
138        for tree_num in tree_nums:
139            # Check if this tree number is under any of the parent tree numbers
140            for parent_tree in parent_trees:
141                if tree_num.startswith(parent_tree + '.'):
142                    descendants.add(heading)
143                    break
144    
145    result = list(descendants)
146    print(f"Found {len(result)} MeSH headings for '{heading}' (including all descendants)")
147    return result