diff --git a/.gitignore b/.gitignore index bbec2b31..5ca2bcbe 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +*.bak +projects example_out example_results hidden_keys.py @@ -15,7 +17,8 @@ EXAMPLE_OUT result_example/ search_terms.md scopus_cache/ - +examples/Termite/01_termite_output +examples/Termite/DOCKERDATA/ # mac results/ poetry.lock diff --git a/CITATION.cff b/CITATION.cff index 758a6a6a..e2f99c26 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -1,4 +1,4 @@ -version: 0.0.43 +version: 0.0.44 message: "If you use this software, please cite it as below." authors: - family-names: Eren @@ -20,7 +20,7 @@ authors: - family-names: Alexandrov given-names: Boian title: "Tensor Extraction of Latent Features (T-ELF)" -version: 0.0.43 +version: 0.0.44 url: https://github.com/lanl/T-ELF doi: 10.5281/zenodo.10257897 date-released: 2023-12-04 diff --git a/README.md b/README.md index 5d8ce9e9..afb6e596 100644 --- a/README.md +++ b/README.md @@ -58,11 +58,11 @@ conda develop . Next, we need to install the optional and additional dependencies. These include optional dependencies for GPU and HPC capabilities, as well as required dependencies like the SpaCy language models. To view all available options, please run: ```shell -python post_install.py --help +telf-post-install --help ``` Install the additional dependencies: ```shell -python post_install.py # use the following, for example, for GPU system: +telf-post-install # use the following, for example, for GPU system: < telf-post-install --gpu> ``` #### Jupyter Setup Tutorial for using the examples ([Link](https://www.maksimeren.com/post/conda-and-jupyter-setup-for-research/)) @@ -129,7 +129,7 @@ python post_install.py # use the following, for example, for GPU system: python=3.11.10 source activate # or use conda activate <...> pip install . -python post_install.py --gpu --hpc-conda +telf-post-install --gpu --hpc-conda ``` ### Darwin @@ -239,5 +239,5 @@ module load miniconda3 conda create --name TELF python=3.11.10 conda activate TELF # or pip install . -python post_install.py --gpu --hpc +telf-post-install --gpu --hpc ``` \ No newline at end of file diff --git a/TELF/applications/Bunny/auto_bunny.py b/TELF/applications/Bunny/auto_bunny.py old mode 100755 new mode 100644 diff --git a/TELF/applications/Lynx/frontend/pages/doc_view.py b/TELF/applications/Lynx/frontend/pages/doc_view.py index 2da74f55..fd760549 100644 --- a/TELF/applications/Lynx/frontend/pages/doc_view.py +++ b/TELF/applications/Lynx/frontend/pages/doc_view.py @@ -181,6 +181,7 @@ for ii in selected_nodes["checked"]: directory = st.session_state.data_map[ii]["path"] peacock_dir = os.path.join(directory, "peacock") + print( peacock_dir ) files = find_files_by_extensions(peacock_dir, extensions=("html", "png")) with st.expander(st.session_state.data_map[ii]["label"]): open_explorer_button(peacock_dir, key=f"button_tab2_{ii}") diff --git a/TELF/applications/Lynx/frontend/pages/helpers/load_project_data.py b/TELF/applications/Lynx/frontend/pages/helpers/load_project_data.py index a39f7643..62ddab11 100644 --- a/TELF/applications/Lynx/frontend/pages/helpers/load_project_data.py +++ b/TELF/applications/Lynx/frontend/pages/helpers/load_project_data.py @@ -9,6 +9,8 @@ import torch import pickle import json +import csv +from collections import defaultdict @st.cache_resource def load_link_data(path): @@ -181,36 +183,213 @@ def add_children(node_name): tree.append(add_children(node["name"])) return tree -def parse_topic_folder_name(folder_name, path=None): +DEBUG = False # prints to terminal + +def _debug(msg): + if DEBUG: + print(msg, flush=True) + +# Matches cluster_for_k=N*.csv (case-insensitive, with optional extra suffix) +_K_FILE_PAT = re.compile(r'^cluster_for_k=(\d+)(?:[^/]*)\.csv$', re.IGNORECASE) + +# Cache per root path +_COUNTS_CACHE = {} +_COUNTS_SRC_CACHE = {} + +def _list_topic_dirs(root_path: str): + try: + dirs = [ + d for d in os.listdir(root_path) + if os.path.isdir(os.path.join(root_path, d)) + ] + # Heuristic: folders that start with digits are topic dirs (e.g., "3-foo", "12", etc.) + topic_dirs = [d for d in dirs if re.match(r'^\d+\b', d)] + return sorted(topic_dirs) + except FileNotFoundError: + return [] + +def _iter_candidate_csvs(root_path: str): + """Yield (abs_path, k) for cluster_for_k=*.csv at root and one level below.""" + if not root_path or not os.path.isdir(root_path): + _debug(f"[scan] Not a dir: {root_path}") + return + + _debug(f"[scan] Looking for cluster_for_k=*.csv under: {root_path}") + + # Top-level files + try: + for entry in os.scandir(root_path): + if entry.is_file(): + m = _K_FILE_PAT.match(entry.name) + if m: + _debug(f"[scan] Found CSV: {entry.name} (k={m.group(1)})") + yield entry.path, int(m.group(1)) + except FileNotFoundError: + pass + + # One level down (subfolders only) + try: + for entry in os.scandir(root_path): + if entry.is_dir(): + for sub in os.scandir(entry.path): + if sub.is_file(): + m = _K_FILE_PAT.match(sub.name) + if m: + _debug(f"[scan] Found CSV (subdir): {entry.name}/{sub.name} (k={m.group(1)})") + yield sub.path, int(m.group(1)) + except FileNotFoundError: + pass + +def _norm_cluster_key(val): + """Normalize cluster key to a stringified integer if possible (e.g., '3', 3, '3.0', 'cluster_3').""" + if val is None: + return None + if isinstance(val, (int, float)): + try: + return str(int(val)) + except Exception: + return str(val).strip() + s = str(val).strip() + m = re.search(r'(\d+)', s) + return m.group(1) if m else s # fallback to raw string + +def _read_counts_map_from_csv(csv_path: str): + """Build counts map {cluster_id(str): count(int)} using the 'cluster' column (case-insensitive).""" + _debug(f"[counts] Reading CSV: {csv_path}") + counts = defaultdict(int) + try: + with open(csv_path, "r", newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + if not reader.fieldnames: + _debug("[counts] No header/fieldnames found.") + return {} + + # find 'cluster' column case-insensitively + cluster_col = None + for name in reader.fieldnames: + if name and name.strip().lower() == "cluster": + cluster_col = name + break + + if not cluster_col: + _debug(f"[counts] 'cluster' column not found in {reader.fieldnames}") + return {} + + row_total = 0 + for row in reader: + row_total += 1 + key = _norm_cluster_key(row.get(cluster_col)) + if key is not None and key != "": + counts[key] += 1 + + _debug(f"[counts] Total rows read: {row_total}. Unique clusters: {len(counts)}") + except Exception as e: + _debug(f"[counts] Error reading {csv_path}: {e}") + return {} + return dict(counts) + +def _choose_best_cluster_csv(root_path: str): + """Choose the best cluster_for_k=*.csv. Prefer k == number of topic dirs; else largest k.""" + candidates = list(_iter_candidate_csvs(root_path)) + if not candidates: + _debug("[choose] No cluster_for_k=*.csv candidates found.") + return None + + topic_dirs = _list_topic_dirs(root_path) + k_target = len(topic_dirs) + _debug(f"[choose] Topic dirs detected: {k_target}") + + # Prefer exact match to number of topic dirs + exact = [p for p in candidates if p[1] == k_target and k_target > 0] + if exact: + # If multiple with same k, pick the shortest path (heuristic) + chosen = sorted(exact, key=lambda x: (len(x[0]), x[0]))[0] + _debug(f"[choose] Using exact k match: {os.path.basename(chosen[0])} (k={chosen[1]})") + return chosen[0] + + # Else pick largest k + candidates.sort(key=lambda x: x[1], reverse=True) + chosen = candidates[0] + _debug(f"[choose] Using largest k: {os.path.basename(chosen[0])} (k={chosen[1]})") + return chosen[0] + +def _get_counts_map(root_path: str): + """Get (and cache) the counts map for this root path.""" + if root_path in _COUNTS_CACHE: + return _COUNTS_CACHE[root_path], _COUNTS_SRC_CACHE.get(root_path) + + csv_path = _choose_best_cluster_csv(root_path) + if not csv_path: + _COUNTS_CACHE[root_path] = {} + _COUNTS_SRC_CACHE[root_path] = None + return {}, None + + counts_map = _read_counts_map_from_csv(csv_path) + _COUNTS_CACHE[root_path] = counts_map + _COUNTS_SRC_CACHE[root_path] = csv_path + return counts_map, csv_path + +def parse_topic_folder_name(folder_name: str, path: str | None = None): """ Extracts topic number, label, and document count from folder name. - - Expected format: topic_number-label_with_underscores-documents_count-documents - Example: '3-label_of_the_topic_9-documents' -> ('3', 'Label of the topic', '9') - - Args: - folder_name (str): Folder name formatted as: -_-documents - - Returns: - tuple: (str, str, str) -> (topic_number, cleaned_label, document_count) - or (None, None, None) if parsing fails. + If the count isn't in the name, read it from cluster_for_k=N.csv at the root (covers all clusters), + using the 'cluster' column to count rows per cluster ID. """ - match = re.match(r'(\d+)-([^-]+)_(\d+)-documents', folder_name) + _debug(f"\n[parse] Folder: {folder_name}") + match = re.match(r'^(\d+)-([^-]+?)(?:[_-](\d+)-documents)?$', folder_name) if match: - topic_number = match.group(1) # Extract topic number - label = match.group(2).replace("_", " ").strip() # Replace underscores with spaces - document_count = match.group(3) # Extract document count + topic_number = match.group(1) + label = match.group(2).replace("_", " ").strip() + document_count = match.group(3) + _debug(f"[parse] Parsed -> number={topic_number}, label='{label}', name_count={document_count}") + + if document_count is None and path: + counts_map, csv_src = _get_counts_map(path) + key = _norm_cluster_key(topic_number) + inferred = counts_map.get(key) + if inferred is not None: + document_count = str(inferred) + _debug(f"[parse] Inferred from CSV ({os.path.basename(csv_src) if csv_src else 'n/a'}): {document_count}") + else: + _debug(f"[parse] No count for cluster={key} in CSV.") + document_count = "Unknown" + + if document_count is None: + document_count = "Unknown" + + _debug(f"[parse] Result -> ({topic_number}, '{label}', {document_count})") return topic_number, label, document_count + elif folder_name.isdigit(): - labels = load_csv_file_items(path, suffix="cluster_summaries", ends=False, column="label") - if path is None or labels is None: - return folder_name, "Topic", "Unknown" - else: - return folder_name, labels[int(folder_name)], "Unknown" + topic_number = folder_name + # Try label mapping (your existing helper) + label = "Topic" + try: + labels = load_csv_file_items(path, suffix="cluster_summaries", ends=False, column="label") + if path is not None and labels is not None: + label = labels[int(folder_name)] + except Exception as e: + _debug(f"[label] Could not map label: {e}") + + document_count = "Unknown" + if path: + counts_map, csv_src = _get_counts_map(path) + key = _norm_cluster_key(topic_number) + inferred = counts_map.get(key) + if inferred is not None: + document_count = str(inferred) + _debug(f"[parse] Inferred (numeric folder) from CSV ({os.path.basename(csv_src) if csv_src else 'n/a'}): {document_count}") + else: + _debug(f"[parse] No count for cluster={key} in CSV (numeric folder).") + + _debug(f"[parse] Numeric folder -> ({topic_number}, '{label}', {document_count})") + return topic_number, label, document_count + _debug("[parse] Unrecognized folder name format.") return None, None, None + def map_folder_to_logical_name(folder_name: str, root_name) -> str: """ Convert a physical folder name (like '*_0', '*_1_2', or '0_1') diff --git a/TELF/applications/Termite/VectorInjector.py b/TELF/applications/Termite/VectorInjector.py new file mode 100644 index 00000000..8675a612 --- /dev/null +++ b/TELF/applications/Termite/VectorInjector.py @@ -0,0 +1,30 @@ +# Termite/VectorInjector.py +import torch +from transformers import AutoTokenizer, AutoModel +from typing import Iterable, List + +class Vectorizer: + """Computes embeddings; does NOT talk to any DB.""" + + def __init__(self, model_name: str = "malteos/scincl", device: str = None, max_length: int = 512): + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + self.model = AutoModel.from_pretrained(model_name) + self.model.eval().to(self.device) + self.max_length = max_length + + @torch.inference_mode() + def encode(self, texts: Iterable[str]) -> List[List[float]]: + outs = [] + for t in texts: + tokens = self.tokenizer( + t if isinstance(t, str) else "", + padding=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt", + ).to(self.device) + hidden = self.model(**tokens).last_hidden_state # [1, L, H] + emb = hidden.mean(dim=1).squeeze(0) # mean pool + outs.append(emb.detach().cpu().tolist()) + return outs diff --git a/TELF/applications/Termite/__init__.py b/TELF/applications/Termite/__init__.py new file mode 100644 index 00000000..32e365f6 --- /dev/null +++ b/TELF/applications/Termite/__init__.py @@ -0,0 +1,5 @@ +# Keep this minimal: just re-export the public Termite class + +from .termite import Termite +from .neo4j_termite.constants import * +__all__ = ["Termite"] diff --git a/TELF/applications/Termite/embedding_store/__init__.py b/TELF/applications/Termite/embedding_store/__init__.py new file mode 100644 index 00000000..6493fb69 --- /dev/null +++ b/TELF/applications/Termite/embedding_store/__init__.py @@ -0,0 +1,19 @@ +# Termite/embedding_store/__init__.py +import os +from .base import EmbeddingStore +from .opensearch_store import OpenSearchStore +from .milvus_store import MilvusStore + +def make_store() -> EmbeddingStore: + backend = os.getenv("EMBEDDING_STORE", "opensearch").lower() + if backend == "milvus": + uri = os.getenv("MILVUS_URI", "http://localhost:19530") + return MilvusStore(uri=uri) + # default: OpenSearch + host = os.getenv("OS_HOST", "localhost") + port = int(os.getenv("OS_PORT", "9200")) + use_ssl = os.getenv("OS_USE_SSL", "false").lower() == "true" + username = os.getenv("OS_USERNAME") + password = os.getenv("OS_PASSWORD") + return OpenSearchStore(host=host, port=port, use_ssl=use_ssl, + username=username, password=password) \ No newline at end of file diff --git a/TELF/applications/Termite/embedding_store/base.py b/TELF/applications/Termite/embedding_store/base.py new file mode 100644 index 00000000..22823b80 --- /dev/null +++ b/TELF/applications/Termite/embedding_store/base.py @@ -0,0 +1,36 @@ +# Termite/embedding_store/base.py +from abc import ABC, abstractmethod +from typing import Iterable, List, Mapping, Optional, Tuple + +class EmbeddingStore(ABC): + """Abstract vector store interface.""" + + @abstractmethod + def ensure_index( + self, + index: str, + dim: int, + metric: str = "cosine", + **kwargs, + ) -> None: + """Create the index/collection if missing.""" + + @abstractmethod + def upsert( + self, + index: str, + ids: List[str], + vectors: List[List[float]], + payloads: Optional[List[Mapping]] = None, + ) -> None: + """Insert or update vectors with optional metadata.""" + + @abstractmethod + def search( + self, + index: str, + query: List[float], + k: int = 5, + **kwargs, + ) -> List[Tuple[str, float, Mapping]]: + """Return [(id, score, payload), ...].""" diff --git a/TELF/applications/Termite/embedding_store/milvus_store.py b/TELF/applications/Termite/embedding_store/milvus_store.py new file mode 100644 index 00000000..b133d1b4 --- /dev/null +++ b/TELF/applications/Termite/embedding_store/milvus_store.py @@ -0,0 +1,60 @@ +# Termite/embedding_store/milvus_store.py +from typing import List, Mapping, Optional, Tuple +from .base import EmbeddingStore + +from pymilvus import MilvusClient, DataType, FieldSchema, CollectionSchema, Collection +import numpy as np + +_METRIC_MAP = {"cosine": "COSINE", "l2": "L2", "dot": "IP"} + +class MilvusStore(EmbeddingStore): + def __init__(self, uri: str = "http://localhost:19530"): + self.client = MilvusClient(uri=uri) + + def ensure_index(self, index: str, dim: int, metric: str = "cosine", **kwargs) -> None: + if self.client.has_collection(index): + return + fields = [ + FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=128, auto_id=False), + FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim), + FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096), + FieldSchema(name="metadata", dtype=DataType.JSON), + ] + schema = CollectionSchema(fields=fields, description="Termite vectors") + self.client.create_collection(collection_name=index, schema=schema, consistency_level="Strong") + self.client.create_index( + collection_name=index, + field_name="embedding", + index_params={"index_type": "HNSW", "metric_type": _METRIC_MAP.get(metric.lower(), "COSINE"), "params": {"M": 16, "efConstruction": 200}}, + ) + + def upsert( + self, + index: str, + ids: List[str], + vectors: List[List[float]], + payloads: Optional[List[Mapping]] = None, + ) -> None: + rows = [] + for i, vid in enumerate(ids): + row = {"id": vid, "embedding": vectors[i]} + if payloads and i < len(payloads) and payloads[i] is not None: + # expect payload to contain "text" and/or "metadata" + row.update(payloads[i]) + rows.append(row) + self.client.insert(collection_name=index, data=rows) + + def search(self, index: str, query: List[float], k: int = 5, **kwargs) -> List[Tuple[str, float, Mapping]]: + res = self.client.search( + collection_name=index, + data=[query], + filter=kwargs.get("filter"), + limit=k, + output_fields=["id", "text", "metadata"], + search_params={"metric_type": "COSINE", "params": {"ef": kwargs.get("ef", 128)}}, + ) + hits = res[0] + out = [] + for h in hits: + out.append((h["entity"]["id"], float(h["distance"]), {"text": h["entity"].get("text"), "metadata": h["entity"].get("metadata")})) + return out diff --git a/TELF/applications/Termite/embedding_store/opensearch_store.py b/TELF/applications/Termite/embedding_store/opensearch_store.py new file mode 100644 index 00000000..7fc8371e --- /dev/null +++ b/TELF/applications/Termite/embedding_store/opensearch_store.py @@ -0,0 +1,164 @@ +# Termite/embedding_store/opensearch_store.py +from typing import List, Mapping, Optional, Tuple, Any +from .base import EmbeddingStore +from opensearchpy import OpenSearch, helpers +import os + +def _normalize_vector(vec: Any) -> List[float]: + """Accept list/tuple/np.ndarray or a string like '[0.1, 0.2]' → list[float].""" + if isinstance(vec, str): + import json, ast + try: + vec = json.loads(vec) + except Exception: + vec = ast.literal_eval(vec) + + try: + import numpy as np + if isinstance(vec, np.ndarray): + vec = vec.tolist() + except Exception: + pass + + if not isinstance(vec, list): + vec = list(vec) + + out = [] + for x in vec: + if x is None: + raise TypeError("Embedding contains None; all values must be floats.") + out.append(float(x)) + if not out: + raise ValueError("Embedding vector is empty.") + return out + +SPACE_MAP = {"cosine": "cosinesimil", "l2": "l2", "dot": "innerproduct"} + +class OpenSearchStore(EmbeddingStore): + def __init__( + self, + host: str = None, + port: int = None, + use_ssl: bool = None, + username: Optional[str] = None, + password: Optional[str] = None, + timeout: int = 60, + ): + # Allow env fallback (useful when Termite sets OS_* env vars) + host = host or os.environ.get("OS_HOST", "localhost") + port = int(port if port is not None else os.environ.get("OS_PORT", "9200")) + use_ssl = bool(str(use_ssl if use_ssl is not None else os.environ.get("OS_USE_SSL", "false")).lower() == "true") + + kwargs = { + "hosts": [{"host": host, "port": port}], + "use_ssl": use_ssl, + "verify_certs": False, # fine for dev + "http_compress": True, + "timeout": timeout, + } + if username and password: + kwargs["http_auth"] = (username, password) + self.client = OpenSearch(**kwargs) + + # ---------- Index management ---------- + def ensure_index( + self, + index: str, + dim: int, + metric: str = "cosine", + ef_search: Optional[int] = 128, # OS 2.13: set at index level (per-query requires >=2.16) + shards: int = 1, + replicas: int = 0, + engine: str = "nmslib", + **kwargs + ) -> None: + if self.client.indices.exists(index=index): + return + + space = SPACE_MAP.get(str(metric).lower(), "cosinesimil") + settings = { + "index": { + "knn": True, + "number_of_shards": int(shards), + "number_of_replicas": int(replicas), + } + } + if ef_search is not None: + # OpenSearch 2.13: tune HNSW candidate breadth here + settings["index"]["knn.algo_param.ef_search"] = int(ef_search) + + body = { + "settings": settings, + "mappings": { + "properties": { + "id": {"type": "keyword"}, + "text": {"type": "text"}, + "metadata": {"type": "object"}, + "embedding": { + "type": "knn_vector", + "dimension": int(dim), + "method": { + "name": "hnsw", + "engine": engine, + "space_type": space, + }, + }, + } + }, + } + self.client.indices.create(index=index, body=body) + + # ---------- Write ---------- + def upsert(self, index, ids, vectors, payloads=None, refresh: bool = True) -> None: + if len(ids) != len(vectors): + raise ValueError("ids and vectors must have the same length.") + actions = [] + for i, vid in enumerate(ids): + src = { + "id": str(vid), + "embedding": _normalize_vector(vectors[i]), + } + if payloads is not None and i < len(payloads) and payloads[i] is not None: + # merge user payloads (e.g., {"text": "...", "metadata": {...}}) + src.update(payloads[i]) + actions.append({"_op_type": "index", "_index": index, "_id": str(vid), "_source": src}) + helpers.bulk(self.client, actions) + if refresh: + self.client.indices.refresh(index=index) + + # ---------- Read ---------- + def search( + self, + index: str, + query, + k: int = 5, + source_fields: Optional[Any] = None, + ) -> List[Tuple[str, float, Mapping]]: + qvec = _normalize_vector(query) + + # _source handling: list or comma-sep string + if source_fields is None: + source_fields = ["id", "text", "metadata"] + elif isinstance(source_fields, str): + source_fields = [s.strip() for s in source_fields.split(",") if s.strip()] + elif not isinstance(source_fields, list): + raise TypeError(f"_source must be list[str] or comma-separated str; got {type(source_fields).__name__}") + + body = { + "size": int(k), + "query": { + "knn": { + # OpenSearch-native syntax for 2.13: + "embedding": {"vector": qvec, "k": int(k)} + } + }, + "_source": source_fields, + } + + resp = self.client.search(index=index, body=body) + hits = resp.get("hits", {}).get("hits", []) + out: List[Tuple[str, float, Mapping]] = [] + for h in hits: + src = h.get("_source", {}) or {} + out.append((h.get("_id", src.get("id")), float(h.get("_score", 0.0)), src)) + return out diff --git a/TELF/applications/Termite/entities/__init__.py b/TELF/applications/Termite/entities/__init__.py new file mode 100644 index 00000000..32392215 --- /dev/null +++ b/TELF/applications/Termite/entities/__init__.py @@ -0,0 +1 @@ +from .return_entities import * \ No newline at end of file diff --git a/TELF/applications/Termite/entities/return_entities.py b/TELF/applications/Termite/entities/return_entities.py new file mode 100644 index 00000000..2e73cc6f --- /dev/null +++ b/TELF/applications/Termite/entities/return_entities.py @@ -0,0 +1,40 @@ +class ReturnEntities: + def __init__(self): + self.all_returns = [] + + def add_ent(self, entity, attributes=None): + """ + Adds an entity to the list of returns. + + Parameters: + ----------- + entity : any + The entity to be added. + attributes : list of tuples, optional + The attributes of the entity. The default is None. + + Returns: + -------- + None + + """ + # Create a dictionary representing the entity and its attributes + entity_dict = { + "ENTITY": entity, + "ATTRIBUTES": attributes if attributes is not None else [] + } + # Append the entity dictionary to the list of returns + return self.all_returns.append(entity_dict) + + + def returns(self): + """ + Returns the list of entities. + + Returns: + -------- + list + A list of dictionaries, where each dictionary contains an entity and its attributes. + """ + # Return the list of entities + return self.all_returns diff --git a/TELF/applications/Termite/misc/material_kg.py b/TELF/applications/Termite/misc/material_kg.py new file mode 100644 index 00000000..343fd18a --- /dev/null +++ b/TELF/applications/Termite/misc/material_kg.py @@ -0,0 +1,45 @@ +# combined_material_clusters_path = "./data/material_injection.csv" +# combined_material_clusters_path = "./affiliation_resolve_materials.csv" +from TELF.applications import Termite +from TELF.applications.Termite.neo4j_termite import * +termite = Termite() +aflow_path = "./material_properties.csv" +# mats_and_NER_path = "./materials_and_MATNER_triplets.csv" +aflow_trip_path = "./aflow_triplets.csv" + +ICSD_TYPE = 'ICSD' +CRYSTAL_SYSTEM_TYPE = 'Crystal_system' +CRYSTAL_CLASS_TYPE ='Crystal_class' +PEARSON_TYPE = 'Pearson' +SPACEGROUP_TYPE = 'Spacegroup' + +MATERIAL_ICSD_RELATION = 'crystal_database_id' +ICSD_SYSTEM_RELATION = 'id_is_system' +ICSD_CLASS_RELATION ='id_is_class' +ICSD_PEARSON_RELATION ='id_is_pearson' +ICSD_SPACEGROUP_RELATION ='id_is_spacegroup' + +material_triplet_map = { + 'ENTITIES': + [ + {ET:MATERIAL_TYPE, FROM_COL: 'material', MAKE_ID_UNIQUE:True}, + {ET:ICSD_TYPE, FROM_COL: 'ICSD', MAKE_ID_UNIQUE:True,ATTR_COL:[{FROM_COL: 'geometry', ATTR_NAME:'geometry'}, {FROM_COL: 'Egap', ATTR_NAME:'Egap'} ]}, + {ET:CRYSTAL_SYSTEM_TYPE, FROM_COL: 'crystal_system', MAKE_ID_UNIQUE:True}, + {ET:CRYSTAL_CLASS_TYPE, FROM_COL: 'crystal_class', MAKE_ID_UNIQUE:True}, + {ET:PEARSON_TYPE, FROM_COL: 'pearson', MAKE_ID_UNIQUE:True}, + {ET:SPACEGROUP_TYPE, FROM_COL: 'spacegroup', MAKE_ID_UNIQUE:True}, + + ], + 'RELATIONS': + [ + {HT:MATERIAL_TYPE, R:MATERIAL_ICSD_RELATION, TT:ICSD_TYPE}, + {HT:ICSD_TYPE, R:ICSD_SYSTEM_RELATION , TT: CRYSTAL_SYSTEM_TYPE, }, + {HT:ICSD_TYPE, R:ICSD_CLASS_RELATION , TT:CRYSTAL_CLASS_TYPE , }, + {HT:ICSD_TYPE, R: ICSD_PEARSON_RELATION, TT: PEARSON_TYPE, }, + {HT:ICSD_TYPE, R: ICSD_SPACEGROUP_RELATION, TT: SPACEGROUP_TYPE, }, + + ] +} +termite.make_unique_constrains( material_triplet_map) +termite.from_csv_to_triplets(aflow_path, aflow_trip_path, material_triplet_map) +termite.update_database_multithreaded(aflow_trip_path,start_from=0,shuffle_rows=True ) \ No newline at end of file diff --git a/TELF/applications/Termite/neo4j_termite/DataInjector.py b/TELF/applications/Termite/neo4j_termite/DataInjector.py new file mode 100644 index 00000000..e0286765 --- /dev/null +++ b/TELF/applications/Termite/neo4j_termite/DataInjector.py @@ -0,0 +1,678 @@ +from .constants import * +import numpy as np, pandas as pd, ast +from neo4j import GraphDatabase +from copy import deepcopy +from tqdm import tqdm +import math +import json +import csv +from concurrent.futures import ThreadPoolExecutor, as_completed + +class InjectorNeo4j: + def __init__(self, + kg_credentials=None, + verbose=False): + """ + Termite, Knowledge graph builder tool. + + Parameters + ---------- + verbose : bool, optional + Verbosity flag. The default is False. + kg_credentials : tuple[str, tuple[str,str]] + first string is url of graph + second tuple is auth, user and password + + Returns + ------- + None. + """ + self.kg_credentials = kg_credentials + self.verbose = verbose + + def setGraphCredentials(self,credentials): + """ + Set the knowledge graph credentials. + + Parameters: + ----------- + credentials : tuple[str, tuple[str,str]] + first string is url of graph + second tuple is auth, user and password + + Returns: + -------- + None + """ + self.kg_credentials = credentials + + def getGraphCredentials(self): + """ + set the knowledge graph credentials + + Parameters: + ----------- + None + + Returns: + -------- + tuple[str, tuple[str,str]] + first string is url of graph + second tuple is auth, user and password + """ + return self.kg_credentials + + + def add_triple(self, + data_container, # Dictionary containing data + head=np.nan, # Head node value + head_type=np.nan, # Head node type + head_attributes=np.nan, # Head node attributes + tail=np.nan, # Tail node value + tail_type=np.nan, # Tail node type + tail_attributes=np.nan, # Tail node attributes + relation=np.nan, # Relation between head and tail + weight=np.nan): # Weight of the relation + """ + Add a triple to the data_container dictionary. + + Parameters: + ----------- + data_container : dict + Dictionary containing data. + head : object, optional + Head node value. The default is np.nan. + head_type : object, optional + Head node type. The default is np.nan. + head_attributes : object, optional + Head node attributes. The default is np.nan. + tail : object, optional + Tail node value. The default is np.nan. + tail_type : object, optional + Tail node type. The default is np.nan. + tail_attributes : object, optional + Tail node attributes. The default is np.nan. + relation : object, optional + Relation between head and tail. The default is np.nan. + weight : object, optional + Weight of the relation. The default is np.nan. + """ + # Append data to data_container + data_container[H].append(head) + data_container[HT].append(head_type) + data_container[HA].append(head_attributes) + data_container[TT].append(tail_type) + data_container[T].append(tail) + data_container[R].append(relation) + data_container[W].append(weight) + data_container[TA].append(tail_attributes) + + def make_unique_constrains(self,column_triplet_map= None, verbose=False, additional_uniques=None): + if not column_triplet_map: + raise ValueError("Need a map for contraints") + + entities = column_triplet_map.get("ENTITIES") + if not entities: + raise ValueError("Need entities in the map") + + if verbose: + print(entities) + + URI, AUTH = self.getGraphCredentials() + with GraphDatabase.driver(URI, auth=AUTH) as driver: + for entity in entities: + make_unique = entity.get(MAKE_ID_UNIQUE) + if make_unique: + Head_type = entity.get(ET) + unique_contraint = f"CREATE CONSTRAINT {Head_type}_id_unique FOR (n:{Head_type}) REQUIRE n.id IS UNIQUE" + try: + + driver.execute_query(unique_contraint, + database_="neo4j", + ) + except Exception as e: + print(f"Failed to create unique constraint on: \n\t{unique_contraint} \n\t {e}") + + def from_csv_to_triplets(self, + csv_path, + save_path, + column_triplet_map=None): + """ + Builds a datafile that maps the raw csv into a head-relation-tail csv + + Parameters: + ----------- + csv_path : str + path to raw data + save_path : str + path to save mapped data + column_triplet_map : dict + Has entities and relations as keys and has a form that is of + the following where the subkeys are defined in the constants file : + {'ENTITIES':[{ ET:TYPE, FROM_COL: COL },], + 'RELATIONS':[{ HT:TYPE, R:TYPE, TT:TYPE, + EXTRACT_H:get_H, EXTRACT_T: get_T },]} + + Returns: + -------- + None + """ + verbose = self.verbose + df = pd.read_csv(csv_path) + + docs = {H:[], R:[], T:[], W:[], HT:[], TT:[], HA:[], TA:[]} + + for index, data in tqdm(df.iterrows(), total = len(df)): + row_entities = {} + # EXTRACT ENTITIES + for entity_map in column_triplet_map['ENTITIES']: + + # Added to unique contraints can be added in the map without prcessing as a column operator + # These entities are expected to have a function to extract them in the relations + if FROM_COL not in entity_map: + continue + + extraction_function = entity_map.get(EXTRACT_ENTITY) + if verbose: + print(entity_map) + + # If there is an extraction function, assume the function will also handle entity attributes + if extraction_function: + args = entity_map.get(ARGS) + if args: + if 'col' in args: # get the column value, remove col name from data + args['data'] = data[args['col']] + del args['col'] + entity = extraction_function(args) + else: + entity = extraction_function() + + + else: + entity = deepcopy(RETURN_TYPE) + if verbose: + print(f"entity_map[FROM_COL] ={entity_map[FROM_COL]}" ) + + if entity_map[FROM_COL] == ROW_INDEX: + entity[ENTITY] = int(index) + if verbose: + print('entity_map[FROM_COL] == ROW_INDEX') + else: + entity[ENTITY] = data[entity_map[FROM_COL]] + + attribute_cols = entity_map.get(ATTR_COL) + if attribute_cols: + for attribute in attribute_cols: + attribute_function = attribute.get(RETREIVAL) + if attribute_function: + # TODO: implement + pass + else: + attr_value = data[attribute[FROM_COL]] + + if type(attr_value) == str or type(attr_value) == int or (type(attr_value) == float and not math.isnan(attr_value)): + if entity[ATTRIBUTES] == None: + entity[ATTRIBUTES] = [(attribute[ATTR_NAME], attr_value)] + else: + entity[ATTRIBUTES].append((attribute[ATTR_NAME], attr_value)) + # else: + # print("Skipped attribute") + + if verbose: + print(f"entity={entity} for entity_map[ET] ={entity_map[ET]}") + + + row_entities[entity_map[ET]] = entity + + if verbose: + print(f"row_entities={row_entities}") + + # EXTRACT RELATIONS + for relation_map in column_triplet_map['RELATIONS']: + triple_details = {H:None, + T:None, + HT: relation_map[HT], + TT: relation_map.get(TT), + HA:None, + TA:None, + W:None, + R: relation_map.get(R), + } + + head_extraction = relation_map.get(EXTRACT_H) + head_args = relation_map.get(ARGS_H, {}) + head_args['data'] = data + + tail_extraction = relation_map.get(EXTRACT_T) + tail_args = relation_map.get(ARGS_T, {}) + tail_args['data'] = data + + # The heads and tails are embedded inline of the data row, and both must be exctracted through the functions passed + if head_extraction and tail_extraction: + head_entities = head_extraction(head_args) # extraction call, pass head type, must return a list of RETURN_TYPE + tail_entities = tail_extraction(tail_args) # extraction call, pass tail type, must return a list of RETURN_TYPE + + pairing = relation_map.get(PAIRING) + if pairing == MANY_TO_MANY or not pairing: # default if no pairing specified + for i, head_entity in enumerate(head_entities): + triple_details[H] = head_entity[ENTITY] + triple_details[W] = head_entity[W] + triple_details[HA] = row_entities[triple_details[HT]][ATTRIBUTES] + + for tail_entity in tail_entities: + if i > 0: + triple_details[HA] = None + triple_details[T] = tail_entity[ENTITY] + triple_details[TA] = tail_entity[ATTRIBUTES] + triple_details[W] = tail_entity[W] + + self.add_triple(docs, **triple_details) + + else: # index pairing + heads_len = len(head_entities) + tails_len = len(tail_entities) + if verbose: + print(head_entities) + print(tail_entities) + print(f"heads_len ={heads_len}, tails_len={tails_len}") + assert heads_len == tails_len + for head_entity, tail_entity in zip(head_entities, tail_entities): + triple_details[H] = head_entity[ENTITY] + triple_details[HA] = head_entity[ATTRIBUTES] + triple_details[W] = head_entity[W] + + triple_details[T] = tail_entity[ENTITY] + triple_details[TA] = tail_entity[ATTRIBUTES] + triple_details[W] = tail_entity[W] + + if triple_details[H] not in [None, 'None'] and triple_details[T] not in [None, 'None']: + self.add_triple(docs, **triple_details) + + # The heads only are embedded inline of the data row, and must be exctracted through the function passed + elif head_extraction: + head_entities = head_extraction(head_args) # extraction call, must return a list of RETURN_TYPE + + if triple_details[TT]: + triple_details[T] = row_entities[triple_details[TT]][ENTITY] + triple_details[TA] = row_entities[triple_details[TT]][ATTRIBUTES] + + + for i, head_entity in enumerate(head_entities): + if i > 0 and triple_details[TT]: + triple_details[TA] = None + + triple_details[H] = head_entity[ENTITY] + triple_details[HA] = head_entity[ATTRIBUTES] + triple_details[W] = head_entity[W] + + if triple_details[H] not in [None, 'None'] and triple_details[T] not in [None, 'None']: + self.add_triple(docs, **triple_details) + + # The tails only are embedded inline of the data row, and must be exctracted through the function passed + elif tail_extraction: + tail_entities = tail_extraction(tail_args ) # extraction call, must return a list of RETURN_TYPE + triple_details[H] = row_entities[triple_details[HT]][ENTITY] # Head is the Row entity looked up through the head type in the details + triple_details[HA] = row_entities[triple_details[HT]][ATTRIBUTES] + + + # print(f'extracting tail for {triple_details[H] }') + + for i , tail_entity in enumerate(tail_entities): + if i > 0: + triple_details[HA] = None + + + triple_details[T] = tail_entity[ENTITY] + triple_details[W] = tail_entity[W] + triple_details[TA] = tail_entity[ATTRIBUTES] + + if triple_details[H] not in [None, 'None'] and triple_details[T] not in [None, 'None']: + # print('Adding triplet') + self.add_triple(docs, **triple_details) + + # The head and tail are both accessed directly through their column maps + else: + if verbose: + print(f"triple_details[HT] = {triple_details[HT]}") + triple_details[H] = row_entities[triple_details[HT]][ENTITY] # Head is the Row entity looked up through the head type in the details + triple_details[HA] = row_entities[triple_details[HT]][ATTRIBUTES] + if triple_details[TT]: + triple_details[T] = row_entities[triple_details[TT]][ENTITY] + triple_details[TA] = row_entities[triple_details[TT]][ATTRIBUTES] + triple_details[W] = row_entities[triple_details[TT]][W] + if verbose: + print(f"triple_details = {triple_details}") + + # if triple_details[H] not in [None, 'None'] and triple_details[T] not in [None, 'None']: + self.add_triple(docs, **triple_details) + + KG_df = pd.DataFrame.from_dict(docs) + KG_df.to_csv(save_path, index=False) + + + def make_attribute_string(self, + attributes, + node_id): + """ + Attribute query contructor for nodes to set the attribuets in the KG + + Parameters: + ----------- + attributes : list + list of tuples containing the attribute name/id and its value + node_id: + which node to assign the attribute + + Returns: + -------- + str + part of a query containing attributes to be appended to a larger query + """ + query_part = '' + attributes = ast.literal_eval(attributes) + if len(attributes): + for attribute in attributes: + attribute_identifier = attribute[0] + attribute_value = attribute[1] + query_part += "SET "+node_id+"." + attribute_identifier + "='" + str(attribute_value).replace('"', '\\"').replace('\'', '\\\'').replace('\\\\\'', '\\\'') + "'" + return query_part + + + def make_triples(self, + driver, + node1, + node1_type, + relation, + node2, + node2_type, + weight, + head_attributes, + tail_attributes, + index = 0): + """ + Parses data into query format for Knowledge graph. Injects query in final step using the driver. + + Parameters: + ----------- + driver : KG driver + to operate the queries + node1 : str + The ID of node 1 + node1_type : str + The graph label for node 1 + relation : str + The edge's graph label + node2 : str + The ID of node + node2_type : str + The graph label for node 2 + weight : float + Weight for the edge + head_attributes : list + Attributes to assign to the head node + tail_attributes : list + Attributes to assign to the tail node + index : int + Current index of Head-Entity-Tail iteration. Helps identify issues in data + + Returns: + -------- + None, int + returns the negative index of queries with problems. No issue, nothing is returned. + """ + # MAKE SURE HEAD NOT NULL + for val in [node1, node1_type]: + if type(val) == float and math.isnan(val): + return -index + + # MAKE ID INT IF NOT STRING + try: + node1 = int(node1) + except: + pass + + # GENERATE HEAD AND ITS ATTRIBUTES + query = "MERGE (a:"+node1_type+" {id: $node1} )" + if type(head_attributes) == str: # not math.isnan(head_attributes): + query += self.make_attribute_string(head_attributes, node_id='a') + + # IF TAIL AND RELATION ARE NOT NULL< GENERATE THEM + skip_relation_tail = False + for val in [relation, node2, node2_type]: + if type(val) == float and math.isnan(val): + skip_relation_tail = True + + if not skip_relation_tail: + try: + node2 = int(node2) + except: + pass + + # GENERATE TAIL AND ITS ATTRIBUTES + query += "MERGE (b:"+node2_type+" {id: $node2} )" + if type(tail_attributes) == str: #not math.isnan(tail_attributes): + query += self.make_attribute_string(tail_attributes, node_id='b') + + # RELATION and optional WEIGHT + if math.isnan(weight): + query += "MERGE (a)-[:"+relation+" ]->(b)" + else: + query += "MERGE (a)-[:"+relation+" {weight: $weight} ]->(b)" + + + # Makes EDGES and NODES in NEO4J + driver.execute_query(query, + node1=node1, node2=node2, weight=weight, + database_="neo4j", + ) + + + def iterate_csv_triplets_into_graph(self, + triplets_path, + start_from =0, + args={}): + """ + Iterates parsed data to call the injection to graph function. + + Parameters: + ----------- + triplets_path : str + Path to the mapped data + args : dict + any additions to be made + + Returns: + -------- + list + Indicies of failed injections + """ + + df = pd.read_csv(triplets_path) + + if start_from: + start = int(len(df)* start_from) + df = df.iloc[start:] + + failed_indicies = [] + + URI, AUTH = self.getGraphCredentials() + with GraphDatabase.driver(URI, auth=AUTH) as driver: + for index, data in tqdm(df.iterrows(), total = len(df)): + + node1 = data[H] + node2 = data[T] + + node1_type = data[HT] + node2_type = data[TT] + + relation = data[R] + weight = data[W] + + head_attributes = data.get(HA, 'nan') + tail_attributes = data.get(TA, 'nan') + index_on_fail = self.make_triples(driver, node1, node1_type, relation, node2, node2_type, weight,head_attributes, tail_attributes, index) + if index_on_fail: + failed_indicies.append(index_on_fail) + + return failed_indicies + + def process_row(self, + data, + index, + driver): + """ + Process a row of data and make triples using the provided driver. + + :param data: The data dictionary containing information about nodes, relations, and attributes. + :param index: The index of the row being processed. + :param driver: The driver object for interacting with the database. + :return: The triples generated by the function call to make_triples. + """ + + node1 = data[H] + node2 = data[T] + + node1_type = data[HT] + node2_type = data[TT] + + relation = data[R] + weight = data[W] + + head_attributes = data.get(HA, 'nan') + tail_attributes = data.get(TA, 'nan') + + return self.make_triples(driver, node1, node1_type, relation, node2, node2_type, weight, head_attributes, tail_attributes, index) + + def update_database_multithreaded(self, + triplets_path, + start_from=0, + shuffle_rows=True): + """ + Iterates parsed data to call the injection to graph function. + + The code does not handle iterations with shuffle_rows=True and start_from != 0. + If shuffle_rows with start_from is needed, manually modify data to accomodate then repass the moddified path. + + Parameters: + ----------- + triplets_path : str + Path to the mapped data + start_from : int + Which row to start from + shuffle_rows : bool + Random shuffle of rows to prevent Neo4j from deadlocking on rows that have common nodes + + Returns: + -------- + list + Indicies of failed injections + """ + df = pd.read_csv(triplets_path) + + if shuffle_rows: + df = df.sample(frac=1).reset_index(drop=True) + df.to_csv(triplets_path) + + failed_indices = [] + + URI, AUTH = self.getGraphCredentials() + + if start_from: + start = int(len(df)* start_from) + print(f"skipping {start}") + df = df.iloc[start:] + + with GraphDatabase.driver(URI, auth=AUTH) as driver: + with ThreadPoolExecutor() as executor: + futures = [executor.submit(self.process_row, data, index, driver) for index, data in df.iterrows()] + # print(len(futures)) + for future in tqdm(as_completed(futures), total=len(futures)): + + index_on_fail = future.result() + if index_on_fail: + failed_indices.append(index_on_fail) + + return failed_indices + + + def write_csv(self, filename, data, headers): + """ + Writes data to a CSV file with specified headers. + + Parameters: + ----------- + filename : str + Path to the CSV file. + data : list of dict + Data to be written to the CSV file. + headers : list of str + Headers for the CSV file. + + Returns: + -------- + None + """ + + # Open the CSV file in 'w' mode with UTF-8 encoding + with open(filename, 'w', newline='', encoding='utf-8') as csvfile: + # Create a DictWriter object from the opened file + writer = csv.DictWriter(csvfile, fieldnames=headers) + + # Write the headers to the CSV file + writer.writeheader() + + # Iterate over the data and write each row to the CSV file + for row in data: + writer.writerow(row) + + def process_json_to_csv(self, + json_path): + """ + Process a JSON file containing nodes and relationships and write + them to CSV files. + + Parameters: + ----------- + json_path : str + Path to the JSON file. + + Returns: + -------- + None + """ + + # Lists to store node and relationship data + nodes_data = [] + relationships_data = [] + + # Open and iterate through the file line by line + with open(json_path, 'r', encoding='utf-8') as file: + for line in file: + try: + # Parse the line as JSON + obj = json.loads(line) + except json.JSONDecodeError: + continue # Skip lines that can't be parsed as JSON + + # If the object is a node, extract its data + if obj['type'] == 'node': + properties = obj.get('properties', {}) + properties[':ID'] = obj.get('id') + properties[':LABEL'] = ';'.join(obj.get('labels', [])) + nodes_data.append(properties) + # If the object is a relationship, extract its data + elif obj['type'] == 'relationship': + properties = obj.get('properties', {}) + properties[':START_ID'] = obj.get('start', {}).get('id') + properties[':END_ID'] = obj.get('end', {}).get('id') + properties[':TYPE'] = obj.get('label') + relationships_data.append(properties) + + # Define CSV headers dynamically based on collected keys + node_headers = list(set(key for row in nodes_data for key in row)) + relationship_headers = list(set(key for row in relationships_data for key in row)) + + # Write node data to a CSV file + self.write_csv('nodes.csv', nodes_data, node_headers) + # Write relationship data to a CSV file + self.write_csv('relationships.csv', relationships_data, relationship_headers) diff --git a/TELF/applications/Termite/neo4j_termite/__init__.py b/TELF/applications/Termite/neo4j_termite/__init__.py new file mode 100644 index 00000000..9737af31 --- /dev/null +++ b/TELF/applications/Termite/neo4j_termite/__init__.py @@ -0,0 +1 @@ +from .constants import * \ No newline at end of file diff --git a/TELF/applications/Termite/neo4j_termite/constants.py b/TELF/applications/Termite/neo4j_termite/constants.py new file mode 100644 index 00000000..b6ef7ee3 --- /dev/null +++ b/TELF/applications/Termite/neo4j_termite/constants.py @@ -0,0 +1,108 @@ +import numpy as np + +""" +General FIELDS +""" +MAKE_ID_UNIQUE = 'unique' + +""" +TRIPLET CSV FIELDS +""" +H = 'head' +HT = 'head_type' +T = 'tail' +TT = 'tail_type' +R = 'relation' +W = 'weight' +HA = 'head_attributes' +TA = 'tail_attributes' +ET = "entity_type" +ROW_INDEX = 'index' +FROM_COL = "from_column" +ATTR_COL = 'attribute_columns' +ATTR_FUNC = 'attribute_function' +ATTR_NAME = 'attribute_name' +ARGS = 'args' +RETREIVAL = 'retrival_operation' +ENTITY = 'entity' +ATTRIBUTES = 'attributes' +EXTRACT_H = 'head_extraction_function' +EXTRACT_T = 'tail_extraction_function' +EXTRACT_ENTITY = 'extract_entity' +PAIRING = "ordering_pairing" +HEAD_TO_MANY = 'head_to_many' +INDEX_PAIRING = 'preserve_index' +MANY_TO_MANY = 'many_to_many' +MANY_TO_TAIL = 'many_to_tail' +ARGS_H = 'head_arguments' +ARGS_T = 'tail_arguments' + + +EMPTY_VALUES = [np.nan, 'None', 'nan'] +RETURN_TYPE = {ENTITY:None, W: None, ATTRIBUTES:None} + +""" +HEADS AND TAILS / NODE TYPING +""" +YEAR_TYPE = 'Year' +TOPIC_TYPE = 'Topic_ID' +COUNTRY_TYPE = 'Country' +KEYWORD_TYPE = 'Keyword' +AUTHOR_ID_TYPE = 'Author_ID' +DOCUMENT_TYPE = 'Document_ID' +DOCUMENT_TYPE_SCOPUS = 'Document_ID_SCOPUS' +SME_WORDS_TYPE = 'SME_word_tag' +NAMED_ENTITY_TYPE = 'Named_Entity' +AFFILIATION_IDENTIFIER_TYPE = 'Affiliation_ID' +PUBLISHER = 'Publisher' +CATEGORY = 'Scopus_category' + +MATERIAL_TYPE = "Material" +AUX_MATERIAL_TYPE = "Auxillary_Material" + + +NER_LOCATION = "NER_location" +NER_PRODUCT = 'NER_product' +NER_ORGANIZATION = 'NER_organization' +NER_GEOPOLITICAL = 'NER_geopolitical_entity' +SME_KEYWORD = 'sme_keyword' +ACRONYM = 'acronym' + + + +""" +RELATION / EDGE TYPES +""" +DOCUMENT_CITES_RELATION = 'cites' +AUTHOR_DOCUMENT_RELATION = 'wrote' +DOCUMENT_CITED_RELATION = 'cited_by' +TOPIC_KEYWORD_RELATION = 'is_about' +DOCUMENT_YEAR_RELATION = 'written_in_year' +DOCUMENT_TOPIC_RELATION = 'is_part_of_topic' +DOCUMENT_AFFILITATION_RELATION = 'is_affiliated_with' +AUTHOR_AFFILITATION_RELATION = 'is_affiliated_with' +AFFILIATION_COUNTRY_RELATION = 'is_in_country' +DOCUMENT_SME_WORD_RELATION = 'has_sme_word' +DOCUMENT_PUBLISHER_RELATION = 'published_by' +DOCUMENT_CATEGORY_RELATION = 'is_in_category' + +DOCUMENT_MATERIAL_RELATION = 'mentions_material' + +MATERIAL_TOPIC_RELATION = 'in_topic' +DOCUMENT_MAIN_MATERIAL_RELATION = 'mentions_main_material' +DOCUMENT_AUX_MATERIAL_RELATION = 'mentions_aux_material' + +TOPIC_AUX_MATERIAL_RELATION = 'contains_aux_material' +DOCUMENT_SME_RELATION='mentions_sme_keyword' +DOCUMENT_LOCATION_RELATION='mentions_location' +DOCUMENT_PRODUCT_RELATION='mentions_product' +DOCUMENT_ORGANIZATION_RELATION='mentions_organization' +DOCUMENT_GEOPOLITICAL_RELATION='mentions_geopolitical_entity' +DOCUMENT_ACRONYM_RELATION='mentions_acronym' + +# WITH ['Topic_ID', 'Keyword', 'Document_ID', 'Document_ID_SCOPUS', 'Affiliation_ID', 'Country', +# 'Year', 'Author_ID', 'SME_word_tag', 'Publisher', 'Scopus_category' ] AS labels +# FOREACH (label IN labels | +# CREATE CONSTRAINT FOR (node:`$label`) REQUIRE node.neo4jImportId IS UNIQUE; +# ) + diff --git a/TELF/applications/Termite/termite.py b/TELF/applications/Termite/termite.py new file mode 100644 index 00000000..2ec30a44 --- /dev/null +++ b/TELF/applications/Termite/termite.py @@ -0,0 +1,317 @@ +# Termite/termite.py +from __future__ import annotations + +import os +from typing import Dict, List, Mapping, Optional, Tuple, Iterable, Any, Union + +from .neo4j_termite.constants import * +from .neo4j_termite.DataInjector import InjectorNeo4j + +# NEW: storage-agnostic vectorizer + store factory +# - Vectorizer computes embeddings only +# - make_store() returns an OpenSearch (default) or Milvus backend +from .VectorInjector import Vectorizer +from .embedding_store import make_store # returns EmbeddingStore + +import pandas as pd + + +class Termite: + """ + Termite, Knowledge graph builder tool. + + Now storage-agnostic for embeddings: + - Default backend: OpenSearch (k-NN) + - Optional backend: Milvus (set EMBEDDING_STORE=milvus or pass embedding_backend='milvus') + + Backward-compat params kept: + - vector_uri: used if backend == 'milvus'; ignored for OpenSearch unless you pass + embedding_backend_config={'host': ..., 'port': ...} + """ + + def __init__( + self, + kg_credentials: Optional[Tuple[str, Tuple[str, str]]] = None, + vector_uri: str = "http://localhost:19530", + db_nme: str = "default", + token: Optional[str] = None, + verbose: bool = False, + embedding_backend: Optional[str] = None, + embedding_backend_config: Optional[Mapping[str, Any]] = None, + model_name: str = "malteos/scincl", + use_gpu: bool = False, + ): + """ + Parameters + ---------- + kg_credentials : tuple[str, tuple[str,str]] + (url, (user, password)) for Neo4j + vector_uri : str + Back-compat. Used as Milvus URI when embedding_backend='milvus'. + db_nme : str + Not used by OpenSearch; kept for compatibility. + token : str | None + Reserved for future auth needs. + verbose : bool + Verbosity flag. + embedding_backend : {'opensearch','milvus'} | None + If None, falls back to ENV EMBEDDING_STORE or 'opensearch'. + embedding_backend_config : Mapping + Backend-specific knobs, e.g. + OpenSearch: {'host':'localhost','port':9200,'use_ssl':False} + Milvus: {'uri':'http://localhost:19530'} + model_name : str + HF model to compute embeddings. + use_gpu : bool + Try to use CUDA for embeddings. + """ + self.verbose = verbose + + # ---- Graph injector (unchanged) ---- + self.graph_injector = InjectorNeo4j(kg_credentials, verbose) + + # ---- Vectorizer (compute embeddings only) ---- + self.vectorizer = Vectorizer(model_name=model_name) + + # ---- Embedding store (pluggable) ---- + # precedence: explicit arg -> ENV -> default('opensearch') + backend = (embedding_backend or os.getenv("EMBEDDING_STORE", "opensearch")).lower() + cfg = dict(embedding_backend_config or {}) + + if backend == "milvus": + # honor legacy vector_uri if not provided in config + cfg.setdefault("uri", vector_uri) + os.environ.setdefault("EMBEDDING_STORE", "milvus") + else: + # default OpenSearch + os.environ.setdefault("EMBEDDING_STORE", "opensearch") + # allow simple overrides + cfg.setdefault("host", os.getenv("OS_HOST", "localhost")) + cfg.setdefault("port", int(os.getenv("OS_PORT", "9200"))) + cfg.setdefault("use_ssl", os.getenv("OS_USE_SSL", "false").lower() == "true") + + # make_store() reads EMBEDDING_STORE + common OS_* / MILVUS_* envs; + # we pass overrides via env for a simple, consolidated factory. + if backend == "milvus": + if "uri" in cfg: + os.environ["MILVUS_URI"] = str(cfg["uri"]) + else: + os.environ["OS_HOST"] = str(cfg["host"]) + os.environ["OS_PORT"] = str(cfg["port"]) + os.environ["OS_USE_SSL"] = "true" if cfg.get("use_ssl") else "false" + + self.store = make_store() + + # book-keeping (optional) + self._default_metric = "cosine" + self._db_name = db_nme + self._token = token + self._model_name = model_name + self._use_gpu = use_gpu + + # ------------------------- + # CSV → triplets (graph) + # ------------------------- + + def from_csv_to_triplets(self, csv_path: str, save_path: str, column_triplet_map: Optional[Mapping] = None): + """ + Builds a datafile that maps the raw csv into a head-relation-tail csv + """ + self.graph_injector.from_csv_to_triplets(csv_path, save_path, column_triplet_map) + + def make_unique_constrains(self, column_triplet_map: Optional[Mapping] = None, verbose: bool = False): + self.graph_injector.make_unique_constrains(column_triplet_map, verbose) + + def iterate_csv_triplets_into_graph(self, triplets_path: str, start_from: int = 0, args: Mapping = {}): + """ + Iterates parsed data to call the injection to graph function. + Returns list of failed indices. + """ + self.graph_injector.iterate_csv_triplets_into_graph(triplets_path, start_from, args) + + def update_database_multithreaded(self, triplets_path: str, start_from: int = 0, shuffle_rows: bool = True): + """ + Iterates parsed data to call the injection to graph function (multithreaded). + """ + self.graph_injector.update_database_multithreaded(triplets_path, start_from, shuffle_rows) + + # ------------------------- + # VECTOR STORE (generic) + # ------------------------- + + def make_vector_schema(self, collection_name: str = "slic", schema: Optional[Mapping] = None, dim: Optional[int] = None, metric: Optional[str] = None): + """ + Creates an index/collection in the selected vector backend. + + For OpenSearch, 'schema' is ignored; we derive a simple mapping with a single + knn_vector field named 'embedding', plus 'id', 'text', 'metadata'. + + Parameters + ---------- + collection_name : str + schema : dict | None + (Kept for compatibility; Milvus users can pass params if they want.) + dim : int | None + Dimension of the embedding. If None, we try to infer from 'schema', + otherwise you must supply it before inserting data. + metric : str | None + 'cosine' (default), 'l2', or 'dot' + """ + metric = (metric or self._default_metric).lower() + if dim is None: + # Try to infer from schema (Milvus-style) else refuse early + if isinstance(schema, dict) and "dim" in schema: + dim = int(schema["dim"]) + else: + raise ValueError("make_vector_schema requires 'dim' (embedding dimension).") + self.store.ensure_index(index=collection_name, dim=dim, metric=metric) + + def inject_vectors(self, collection_name: str, data: List[Mapping[str, Any]]): + """ + Insert / upsert vectors into the selected backend. + + Expected 'data' shape (per item): + { + 'id': str, + 'embedding': List[float], # length = dim + # optional user payload: + 'text': str, + 'metadata': Mapping[str, Any] + } + """ + ids: List[str] = [] + vecs: List[List[float]] = [] + payloads: List[Mapping[str, Any]] = [] + + for row in data: + if "id" not in row or "embedding" not in row: + raise ValueError("Each data row must contain 'id' and 'embedding'.") + ids.append(str(row["id"])) + vecs.append(list(row["embedding"])) + payloads.append({k: v for k, v in row.items() if k not in ("id", "embedding")}) + + self.store.upsert(index=collection_name, ids=ids, vectors=vecs, payloads=payloads) + + # ------------------------- + # HELPERS to shape data + # ------------------------- + + def df_to_data( + self, + embeddings: List[List[float]], + df_path: str, + columns_collection_map: Mapping[str, str], + id_col: Optional[str] = None, + text_col: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """ + Converts a CSV into a list of upsertable vector rows. + + Parameters + ---------- + embeddings : list[list[float]] + Embeddings in the same row-order as the CSV. + df_path : str + Path to CSV. + columns_collection_map : dict + Mapping of wanted CSV columns → payload keys (e.g., {"abstract":"text","paper_id":"id"}). + If 'id' or 'text' are not provided here, we will fallback to id_col/text_col args. + id_col, text_col : Optional[str] + Fallbacks for the identifier and text payload. + + Returns + ------- + List[dict] + Each dict at least has {'id', 'embedding'} plus any mapped payload fields. + """ + df = pd.read_csv(df_path) + if len(df) != len(embeddings): + raise ValueError(f"embeddings length ({len(embeddings)}) != rows in CSV ({len(df)})") + + # Determine id/text columns + mapped_id = next((src for src, dst in columns_collection_map.items() if dst == "id"), None) or id_col + mapped_text = next((src for src, dst in columns_collection_map.items() if dst == "text"), None) or text_col + + data: List[Dict[str, Any]] = [] + for i, row in df.iterrows(): + payload: Dict[str, Any] = {} + + # Map requested columns + for src, dst in columns_collection_map.items(): + if src in df.columns: + payload[dst] = row[src] + + # Ensure id/text presence if desired + rid = str(row[mapped_id]) if mapped_id and mapped_id in df.columns else str(i) + if "id" not in payload: + payload["id"] = rid + + if mapped_text and mapped_text in df.columns and "text" not in payload: + payload["text"] = row[mapped_text] + + data.append({"id": payload.pop("id"), "embedding": embeddings[i], **payload}) + + return data + + # ------------------------- + # EMBEDDING COMPUTATION + # ------------------------- + + def compute_embeddings( + self, + df: "pd.DataFrame", + model_name: str = "SCINCL", + use_gpu: bool = False, + text_column: Optional[str] = None, + ) -> Dict[Any, List[float]]: + """ + Computes embeddings for a pandas DataFrame. + + Assumes there is a textual column to embed. If `text_column` is None, + we try 'text', then the first object/string dtype column. + + Returns + ------- + dict: {row_index -> embedding_vector} + """ + # choose column + col = text_column + if col is None: + if "text" in df.columns: + col = "text" + else: + # pick first object-dtype column as a best-effort fallback + obj_cols = [c for c in df.columns if df[c].dtype == object] + if not obj_cols: + raise ValueError("No text column found. Provide 'text_column'.") + col = obj_cols[0] + + texts: List[str] = df[col].astype(str).tolist() + + # swap model if caller passed one + if model_name and model_name != self._model_name: + self.vectorizer = Vectorizer(model_name=model_name) + self._model_name = model_name + + # note: Vectorizer selects CUDA automatically if available; we keep the flag for API parity + embs = self.vectorizer.encode(texts) + + # map back to index + return {idx: embs[i] for i, idx in enumerate(df.index.tolist())} + + # ------------------------- + # OPTIONAL: convenience search + # ------------------------- + + def search_vectors( + self, + collection_name: str, + query_vec: List[float], + k: int = 5, + **kwargs, + ) -> List[Tuple[str, float, Mapping]]: + """ + Simple k-NN search wrapper to the active backend. + Returns a list of (id, score, payload). + """ + return self.store.search(index=collection_name, query=query_vec, k=k, **kwargs) diff --git a/TELF/applications/__init__.py b/TELF/applications/__init__.py index 9d88e0ec..764a6383 100644 --- a/TELF/applications/__init__.py +++ b/TELF/applications/__init__.py @@ -2,6 +2,7 @@ sys.path += ["Cheetah"] sys.path += ["Bunny"] sys.path += ["Penguin"] +sys.path += ["Termite"] # Cheetah @@ -16,3 +17,8 @@ # Penguin from .Penguin.penguin import Penguin + +# Termite +from .Termite.termite import Termite +from .Termite.neo4j_termite.constants import * + diff --git a/TELF/factorization/HNMFk.py b/TELF/factorization/HNMFk.py index 14d5040e..22efa6cc 100644 --- a/TELF/factorization/HNMFk.py +++ b/TELF/factorization/HNMFk.py @@ -609,21 +609,22 @@ def _process_node(self, Ks, # check if leaf node status based on number of samples # if (current_node.num_samples == 1): - current_node.leaf = True - current_node.exception = True - pickle_path = os.path.join(str(node_save_path), f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} - + return self._finalize_leaf( + current_node, node_save_path, + reason="single-sample cluster", + details={"num_samples": current_node.num_samples}, + mark_exception=True + ) # # Sample threshold check for leaf node determination # if self.sample_thresh > 0 and (current_node.num_samples <= self.sample_thresh): - current_node.leaf = True - pickle_path = os.path.join(f'{node_save_path}', f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} + return self._finalize_leaf( + current_node, node_save_path, + reason="below sample threshold", + details={"num_samples": current_node.num_samples, "sample_thresh": self.sample_thresh} + ) # # obtain the current X @@ -644,12 +645,13 @@ def _process_node(self, Ks, # Based on number of features or samples, no seperation possible # if min(curr_X.shape) <= 1: - current_node.leaf = True - current_node.exception = True - pickle_path = os.path.join(f'{node_save_path}', f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} - + return self._finalize_leaf( + current_node, node_save_path, + reason="matrix too small to factorize", + details={"X_shape": tuple(curr_X.shape)}, + mark_exception=True + ) + # # prepare the current nmfk parameters # @@ -665,11 +667,12 @@ def _process_node(self, Ks, # Ks = self._adjust_curr_Ks(curr_X.shape, Ks) if len(Ks) == 0 or (len(Ks) == 1 and Ks[0] < 2): - current_node.leaf = True - current_node.exception = True - pickle_path = os.path.join(str(node_save_path), f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} + return self._finalize_leaf( + current_node, node_save_path, + reason="no valid K values for further split", + details={"Ks": list(Ks), "X_shape": tuple(curr_X.shape)}, + mark_exception=True + ) # # apply nmfk @@ -681,11 +684,12 @@ def _process_node(self, Ks, # Check if decomposition was not possible # if results is None: - current_node.leaf = True - current_node.exception = True - pickle_path = os.path.join(str(node_save_path), f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} + return self._finalize_leaf( + current_node, node_save_path, + reason="NMFk failed to decompose", + details={"node_depth": current_node.depth, "num_samples": current_node.num_samples}, + mark_exception=True + ) # # latent factors @@ -723,10 +727,24 @@ def _process_node(self, Ks, # leaf node based on depth limit or single cluster or all samples in same cluster if ((current_node.depth >= self.depth) and self.depth > 0) or current_node.k == 1 or n_clusters == 1: - current_node.leaf = True - pickle_path = os.path.join(f'{node_save_path}', f'node_{current_node.node_name}.p') - pickle.dump(current_node, open(pickle_path, "wb")) - return {"name":node_name, "target_jobs":[], "node_save_path":pickle_path} + stop_details = { + "depth": current_node.depth, + "depth_limit": self.depth, + "predicted_k": current_node.k, + "n_clusters_observed": n_clusters + } + # Pick a specific reason to make the text friendlier: + if (current_node.depth >= self.depth) and self.depth > 0: + reason = "depth limit reached" + elif current_node.k == 1: + reason = "predicted k == 1 (no further structure)" + elif n_clusters == 1: + reason = "all samples assigned to one cluster" + else: + reason = "stopping condition met" + + return self._finalize_leaf(current_node, node_save_path, reason=reason, details=stop_details) + # # go through each topic/cluster @@ -1110,4 +1128,56 @@ def _resolve_path(self, original_node_path: str): # 2. Combine it with the new experiment base new_node_path = new_experiment_base / relative_path - return new_node_path \ No newline at end of file + return new_node_path + + + """ + DEBUG FUNCTIONS + ---------------- + """ + def _write_stop_reason(self, node_save_path: str, reason: str, details: dict | None = None): + """ + Write a plain-text explanation for why decomposition did not proceed. + Creates /stop_reason.txt next to the node pickle. + """ + try: + reason_path = os.path.join(node_save_path, "stop_reason.txt") + with open(reason_path, "w") as f: + f.write(f"reason: {reason}\n") + if details: + for k, v in details.items(): + f.write(f"{k}: {v}\n") + except Exception as e: + # Don't break the run just because we couldn't write the note + warnings.warn(f"Could not write stop_reason.txt at {node_save_path}: {e}") + + def _finalize_leaf(self, current_node, node_save_path: str, reason: str, details: dict | None = None, *, + mark_exception: bool = False): + """ + Common exit path for any condition that stops further decomposition. + - marks leaf/exception on the node + - persists the node + - writes stop_reason.txt with a friendly explanation + - returns the scheduler payload + """ + current_node.leaf = True + current_node.exception = bool(mark_exception) + pickle_path = os.path.join(str(node_save_path), f'node_{current_node.node_name}.p') + + # persist the node + pickle.dump(current_node, open(pickle_path, "wb")) + + # write a human-readable reason file + # self._write_stop_reason(node_save_path, reason, details) + + + details = details or {} + details.update({ + "cluster_on": self.cluster_on, + "num_samples": current_node.num_samples, + "len_original_indices": len(current_node.original_indices), + }) + details.update(getattr(current_node, "user_node_data", {}) or {}) + self._write_stop_reason(node_save_path, reason, details) + + return {"name": current_node.node_name, "target_jobs": [], "node_save_path": pickle_path} diff --git a/TELF/helpers/figures.py b/TELF/helpers/figures.py index 671c42f9..75cb5334 100644 --- a/TELF/helpers/figures.py +++ b/TELF/helpers/figures.py @@ -27,6 +27,27 @@ def plot_authors_graph(df, id_col='s2_author_ids', name_col='s2_authors', title='Co-Authors Graph', width=900, height=900, max_node_size=50, min_node_size=3): G = create_authors_graph(df, id_col) + + # -------------------------- bail out gracefully on empty graphs -------------------------- + if G.number_of_nodes() == 0 or G.number_of_edges() == 0: + fig = go.Figure() + fig.update_layout( + title=f"{title} — no data to display", + width=width, + height=height, + xaxis=dict(visible=False), + yaxis=dict(visible=False), + showlegend=False, + margin=dict(l=20, r=20, t=60, b=20), + ) + # Optional: a centered annotation so the HTML isn't just a blank box + fig.add_annotation( + text="No relationships found for the given data/columns.", + x=0.5, y=0.5, xref="paper", yref="paper", + showarrow=False + ) + return fig + pos = nx.spring_layout(G) # position nodes using networkx's spring layout name_map = get_id_to_name(df, name_col, id_col) diff --git a/TELF/pipeline/__init__.py b/TELF/pipeline/__init__.py index f0f242ba..ba9a6ec5 100644 --- a/TELF/pipeline/__init__.py +++ b/TELF/pipeline/__init__.py @@ -65,3 +65,12 @@ from .blocks.ocelot_filter_block import OcelotFilterBlock from .blocks.auto_bunny_simple_block import AutoBunnySimpleBlock from .blocks.term_table_block import TermTableBlock +from .blocks.spacey_NER_block import SpacyNERBlock + + +from .blocks.collect_hnmfk_leaf_block import CollectHNMFkLeafBlock +from .blocks.termite_neo4j_block import TermiteNeo4jBlock +from .blocks.termite_vector_block import TermiteVectorBlock + +from .blocks.author_affiliation_tables import AffiliationsAndAuthorsBlock +from .blocks.block_helpers.KernelServer import KernelTiedServer diff --git a/TELF/pipeline/block_manager.py b/TELF/pipeline/block_manager.py index 032a304b..6c381274 100644 --- a/TELF/pipeline/block_manager.py +++ b/TELF/pipeline/block_manager.py @@ -57,6 +57,48 @@ def __init__( self.describe_io() def __call__(self) -> DataBundle: + base = Path(self.bundle[SAVE_DIR_BUNDLE_KEY]) + + # 1) Preflight: rename block directories and update checkpoint paths on disk + # MUST return: + # - base_to_display: {"SemanticHNMFk": "06_SemanticHNMFk", ...} + # - prefix_map: {"/.../07_SemanticHNMFk": "/.../06_SemanticHNMFk", ...} + base_to_display, prefix_map = _renumber_dirs_and_update_ckpts(base, self.blocks) + + # 2) Update in-memory objects to new prefixes BEFORE any block runs + # 2a) Rewrite any stored paths inside the bundle values + try: + for base_key, bucket in list(self.bundle._store.items()): + for tag, val in list(bucket.items()): + if tag == "_latest": + continue + # Apply all variants + new_val = val + for old_pref, new_pref in prefix_map.items(): + new_val = _deep_replace_in_obj(new_val, {old_pref: new_pref}) + bucket[tag] = new_val + except Exception: + pass + + # Set display tags and rewrite block settings (init/call) in memory + for block in self.blocks: + bt = _base_tag(getattr(block, "_original_tag", block.tag)) + if not hasattr(block, "_original_tag"): + block._original_tag = bt + block.tag = base_to_display.get(bt, bt) + + if isinstance(getattr(block, "init_settings", None), dict): + new_init = block.init_settings + for old_pref, new_pref in prefix_map.items(): + new_init = _deep_replace_in_obj(new_init, {old_pref: new_pref}) + block.init_settings = new_init + + if isinstance(getattr(block, "call_settings", None), dict): + new_call = block.call_settings + for old_pref, new_pref in prefix_map.items(): + new_call = _deep_replace_in_obj(new_call, {old_pref: new_pref}) + block.call_settings = new_call + total = len(self.blocks) log_dir: Path | None = None progress_fp = None @@ -69,6 +111,7 @@ def __call__(self) -> DataBundle: progress_fp.write("# IO table\n" + "\n".join(table_lines) + "\n\n") progress_fp.flush() + # 3) Run blocks for idx, block in enumerate(self.blocks, 1): # Override the block’s load_checkpoint flag if requested if self.force_checkpoint is not None: @@ -89,7 +132,6 @@ def __call__(self) -> DataBundle: buf_err.write(f"⚠️ Exception in block {block.tag}:\n") traceback.print_exc(file=buf_err) captured = buf_out.getvalue() + buf_err.getvalue() - # Write error log and re-raise ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S") (log_dir / f"{idx:02d}_{block.tag}_{ts}.log").write_text(captured, encoding="utf-8") raise @@ -108,6 +150,16 @@ def __call__(self) -> DataBundle: traceback.print_exc() raise + # 4) Alias outputs: also store under the base tag (no NN_ prefix) + disp_tag = block.tag # e.g., "06_SemanticHNMFk" + base_tag = _base_tag(disp_tag) # e.g., "SemanticHNMFk" + for base_key in self.bundle.keys_by_tag(disp_tag): + try: + val = self.bundle[f"{disp_tag}.{base_key}"] + self.bundle[f"{base_tag}.{base_key}"] = val + except KeyError: + pass + elapsed = time.perf_counter() - t0 if self.progress: print(f"✓ [{idx}/{total}] {block.tag} finished in {elapsed:,.2f}s") @@ -121,6 +173,7 @@ def __call__(self) -> DataBundle: progress_fp.close() return self.bundle + # ------------------------------------------------------------------ # # helper – produce describe-io lines without printing # # ------------------------------------------------------------------ # @@ -359,7 +412,7 @@ def save_settings(self) -> None: for idx, blk in enumerate(self.blocks): fn = saved_dir / f"{idx}_{blk.__class__.__name__}.json" - fn.write_text(jsonpickle.encode(blk), encoding="utf-8") + fn.write_text(jsonpickle.encode(blk, keys=True), encoding="utf-8") def load_saved_settings(self) -> None: """ @@ -402,3 +455,259 @@ def _ensure_result_path(self) -> None: """ if "result_path" not in self.bundle: self.bundle["result_path"] = Path.cwd() / "results" + + + +import json, os, re, uuid, pickle +from pathlib import Path +from typing import Any + +INDEXED_DIR_RE = re.compile(r"^\d+_") +TEXT_EXTS = {".json", ".txt", ".yaml", ".yml", ".ini", ".cfg", ".csv", ".tsv"} +PICKLE_EXTS = {".p", ".pkl", ".pickle"} + +def _base_tag(name: str) -> str: + return INDEXED_DIR_RE.sub("", name) + +def _display_name(tag: str, idx: int, width: int) -> str: + return f"{idx:0{width}d}_{tag}" + +def _find_existing_dir(base: Path, tag: str) -> Path | None: + candidates = sorted( + base.glob(f"[0-9][0-9]*_{tag}"), + key=lambda p: p.stat().st_mtime if p.exists() else 0, + reverse=True, + ) + if candidates: + return candidates[0] + plain = base / tag + return plain if plain.exists() else None + +def _collect_ckpt_files(base: Path) -> list[Path]: + return list(base.rglob("__checkpoints__.json")) + +def _rewrite_ckpt_paths(ckpt_file: Path, prefix_map: dict[str, str]) -> bool: + try: + data = json.loads(ckpt_file.read_text(encoding="utf-8")) + except Exception: + return False + changed = False + for k, v in list(data.items()): + if isinstance(v, str): + for old_prefix, new_prefix in prefix_map.items(): + if v == old_prefix or v.startswith(old_prefix + os.sep): + data[k] = new_prefix + v[len(old_prefix):] + changed = True + if changed: + ckpt_file.write_text(json.dumps(data, indent=2), encoding="utf-8") + return changed + +def _deep_replace(obj: Any, old_prefix: str, new_prefix: str, seen: set[int] | None = None) -> tuple[bool, Any]: + if seen is None: + seen = set() + oid = id(obj) + if oid in seen: + return False, obj + seen.add(oid) + + if isinstance(obj, str): + if old_prefix in obj: + return True, obj.replace(old_prefix, new_prefix) + return False, obj + + if isinstance(obj, dict): + changed = False + out = {} + for k, v in obj.items(): + ck, nk = _deep_replace(k, old_prefix, new_prefix, seen) if isinstance(k, str) else (False, k) + cv, nv = _deep_replace(v, old_prefix, new_prefix, seen) + changed = changed or ck or cv + out[nk] = nv + return changed, out + + if isinstance(obj, list): + changed = False + out = [] + for v in obj: + cv, nv = _deep_replace(v, old_prefix, new_prefix, seen) + changed = changed or cv + out.append(nv) + return changed, out + + if isinstance(obj, tuple): + changed = False + out_list = [] + for v in obj: + cv, nv = _deep_replace(v, old_prefix, new_prefix, seen) + changed = changed or cv + out_list.append(nv) + return changed, tuple(out_list) + + try: + attrs = vars(obj) + except Exception: + return False, obj + + changed = False + for k, v in list(attrs.items()): + cv, nv = _deep_replace(v, old_prefix, new_prefix, seen) + if cv: + try: + setattr(obj, k, nv) + changed = True + except Exception: + pass + return changed, obj + +def _rewrite_internal_paths_in_tree(root: Path, prefix_map: dict[str, str]) -> None: + """ + Walk files under `root` and replace occurrences of *any* old→new prefix. + Handles text and pickle files; best-effort, silent on failures. + """ + for p in root.rglob("*"): + if not p.is_file(): + continue + ext = p.suffix.lower() + + if ext in TEXT_EXTS: + try: + s = p.read_text(encoding="utf-8", errors="ignore") + changed = False + for old_prefix, new_prefix in prefix_map.items(): + if old_prefix in s: + s = s.replace(old_prefix, new_prefix) + changed = True + if changed: + p.write_text(s, encoding="utf-8") + except Exception: + pass + continue + + if ext in PICKLE_EXTS: + try: + with p.open("rb") as f: + obj = pickle.load(f) + changed_any = False + for old_prefix, new_prefix in prefix_map.items(): + changed, obj = _deep_replace(obj, old_prefix, new_prefix) + changed_any = changed_any or changed + if changed_any: + with p.open("wb") as f: + pickle.dump(obj, f) + except Exception: + pass + continue + +def _renumber_dirs_and_update_ckpts(base: Path, blocks: list) -> tuple[dict[str, str], dict[str, str]]: + """ + Rename block directories to NN_, rewrite checkpoint JSONs, + and migrate internal absolute/relative paths inside files under renamed dirs. + RETURNS: + (base_to_display, prefix_map) where prefix_map includes abs+rel variants. + """ + base.mkdir(parents=True, exist_ok=True) + + width = max(2, len(str(len(blocks)))) + desired: list[tuple[str, str]] = [] + for i, block in enumerate(blocks, start=1): + bt = _base_tag(getattr(block, "_original_tag", block.tag)) + desired.append((bt, _display_name(bt, i, width))) + + # Plan renames + rename_plan: list[tuple[Path, Path]] = [] + for (bt, new_disp) in desired: + src = _find_existing_dir(base, bt) + if src is None: + continue + dst = base / new_disp + if src.resolve() != dst.resolve(): + rename_plan.append((src, dst)) + + # Two-phase rename to avoid collisions + temp_map: dict[Path, Path] = {} + for src, dst in rename_plan: + if not src.exists(): + continue + tmp = base / (src.name + f".tmp-{uuid.uuid4().hex[:8]}") + src.rename(tmp) + temp_map[tmp] = dst + + # Move temps to final + for tmp, dst in temp_map.items(): + if dst.exists(): + if dst.is_dir() and tmp.is_dir(): + for item in tmp.iterdir(): + target = dst / item.name + if not target.exists(): + item.rename(target) + tmp.rmdir() + else: + n = 1 + alt = Path(str(dst) + f".old{n}") + while alt.exists(): + n += 1 + alt = Path(str(dst) + f".old{n}") + tmp.rename(alt) + else: + tmp.rename(dst) + + # Build a rich prefix map (absolute + relative variants) + prefix_map = _build_prefix_map_variants(base, rename_plan) + + # Rewrite checkpoint JSON files using all variants + for ckpt in _collect_ckpt_files(base): + _rewrite_ckpt_paths(ckpt, prefix_map) + + # Rewrite internals inside each renamed directory tree (text+pickle) using all variants + for src, dst in rename_plan: + try: + _rewrite_internal_paths_in_tree(dst, prefix_map) + except Exception: + pass + + base_to_display = {bt: new_disp for (bt, new_disp) in desired} + return base_to_display, prefix_map + + +def _deep_replace_in_obj(obj, prefix_map: dict[str, str]): + from pathlib import Path as _Path + if isinstance(obj, (str, _Path)): + s = str(obj) + for old, new in prefix_map.items(): + if s == old or s.startswith(old + os.sep): + s = new + s[len(old):] + return _Path(s) if isinstance(obj, _Path) else s + if isinstance(obj, dict): + return { _deep_replace_in_obj(k, prefix_map) if isinstance(k, (str, _Path)) else k: + _deep_replace_in_obj(v, prefix_map) for k, v in obj.items() } + if isinstance(obj, list): + return [ _deep_replace_in_obj(v, prefix_map) for v in obj ] + if isinstance(obj, tuple): + return tuple(_deep_replace_in_obj(v, prefix_map) for v in obj) + return obj + +def _build_prefix_map_variants(base: Path, rename_plan: list[tuple[Path, Path]]) -> dict[str, str]: + """ + For each (src,dst) directory rename, return a mapping that includes: + - absolute: /abs/.../07_Tag -> /abs/.../06_Tag + - relative (from CWD): src -> dst (e.g., 'example_results/.../07_Tag' -> '.../06_Tag') + This catches both absolute and relative paths embedded in files or memory. + """ + m: dict[str, str] = {} + for src, dst in rename_plan: + # absolute variants + abs_old = str(src.resolve()) + abs_new = str(dst.resolve()) + m[abs_old] = abs_new + + # relative variants (as written on disk; your code uses Path(...) directly) + rel_old = str(src) # typically 'example_results/.../07_Tag' + rel_new = str(dst) + m[rel_old] = rel_new + + # Sometimes code stores a trailing slash; add those too + if not abs_old.endswith(os.sep): + m[abs_old + os.sep] = abs_new + os.sep + if not rel_old.endswith(os.sep): + m[rel_old + os.sep] = rel_new + os.sep + return m diff --git a/TELF/pipeline/blocks/__init__.py b/TELF/pipeline/blocks/__init__.py index ff6eb888..99920a63 100644 --- a/TELF/pipeline/blocks/__init__.py +++ b/TELF/pipeline/blocks/__init__.py @@ -58,4 +58,11 @@ from .peacock_stats_block import PeacockStatsBlock from .ocelot_filter_block import OcelotFilterBlock from .auto_bunny_simple_block import AutoBunnySimpleBlock -from .term_table_block import TermTableBlock \ No newline at end of file +from .term_table_block import TermTableBlock +from .spacey_NER_block import SpacyNERBlock +from .collect_hnmfk_leaf_block import CollectHNMFkLeafBlock +from .termite_neo4j_block import TermiteNeo4jBlock +from .termite_vector_block import TermiteVectorBlock +from .author_affiliation_tables import AffiliationsAndAuthorsBlock + +from .block_helpers.KernelServer import KernelTiedServer \ No newline at end of file diff --git a/TELF/pipeline/blocks/artic_fox_block.py b/TELF/pipeline/blocks/artic_fox_block.py index 160915f5..39119072 100644 --- a/TELF/pipeline/blocks/artic_fox_block.py +++ b/TELF/pipeline/blocks/artic_fox_block.py @@ -1,14 +1,21 @@ from pathlib import Path -from typing import Dict, Sequence, Any, Tuple +from typing import Dict, Any from .base_block import AnimalBlock -from .data_bundle import DataBundle +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY from ...post_processing import ArcticFox from ...factorization import HNMFk + class ArticFoxBlock(AnimalBlock): + """ + Block wrapper for the ArcticFox post-process/label/stats pipeline. + Use call_settings['steps'] to run any subset: ["post"], ["label"], ["stats"], + or combinations like ["post","label"], ["label","stats"], ["post","stats"], ["post","label","stats"]. + If steps is None, legacy behavior uses the label_clusters/generate_stats booleans. + """ - CANONICAL_NEEDS = ("df", 'vocabulary', "model_path",) + CANONICAL_NEEDS = ("df", "vocabulary", "model_path") def __init__( self, @@ -21,31 +28,29 @@ def __init__( call_settings: Dict[str, Any] = None, **kw, ) -> None: - + self.col = col default_init = { - 'clean_cols_name': self.col, - 'embedding_model': "SCINCL", + "clean_cols_name": self.col, + "embedding_model": "SCINCL", } default_call = { - 'ollama_model': "llama3.2:3b-instruct-fp16", # Language model used for semantic label generation - 'label_clusters': True, # Enable automatic labeling of clusters - 'generate_stats': True, # Generate cluster-level statistics - 'process_parents': True, # Propagate labels or stats upward through the hierarchy - 'skip_completed': True, # Skip processing of nodes already labeled/stored - 'label_criteria': { # Rules to filter generated labels - "minimum words": 2, - "maximum words": 6 - }, - 'label_info': { # Additional metadata to associate with generated labels - "source": "Science" - }, - 'number_of_labels': 5 # Number of candidate labels to generate per node + "ollama_model": "llama3.2:3b-instruct-fp16", # Language model used for semantic label generation + "label_clusters": True, # Back-compat: used when steps is None + "generate_stats": True, # Back-compat: used when steps is None + "process_parents": True, + "skip_completed": True, + "label_criteria": {"minimum words": 2, "maximum words": 6}, + "label_info": {"source": "Science"}, + "number_of_labels": 5, + # NEW: choose subset explicitly; None keeps legacy boolean behavior + # Examples: ["post"], ["label"], ["stats"], ["post","label"], ["label","stats"], ["post","stats"], ["post","label","stats"] + "steps": None, } super().__init__( - needs = needs, - provides = provides, + needs=needs, + provides=provides, init_settings=self._merge(default_init, init_settings), call_settings=self._merge(default_call, call_settings), tag=tag, @@ -53,15 +58,40 @@ def __init__( ) def run(self, bundle: DataBundle) -> None: - df = self.load_path(bundle[self.needs[0]]) + # Resolve inputs + df = self.load_path(bundle[self.needs[0]]) vocabulary = self.load_path(bundle[self.needs[1]]) - model = HNMFk(experiment_name=bundle[self.needs[2]]) - model.load_model() # Loads model from the provided experiment_name path - pipeline = ArcticFox( - model=model, - **self.init_settings + raw_model_path = str(bundle[self.needs[2]]) + + try: + resolved_model_path = str(Path(raw_model_path).expanduser().resolve()) + except Exception: + resolved_model_path = raw_model_path + + # Load HNMFk model + model = HNMFk(experiment_name=raw_model_path) + model.load_model() + + # Run selected steps (order enforced inside ArcticFox) + pipeline = ArcticFox(model=model, **self.init_settings) + pipeline.run_full_pipeline( + data_df=df, + vocab=vocabulary, + **self.call_settings ) - pipeline.run_full_pipeline(data_df = df, - vocab = vocabulary, - **self.call_settings) - bundle[f"{self.tag}.{self.provides[0]}"] = "Done" \ No newline at end of file + + # Write a lightweight status checkpoint + status_value = "Done" + if SAVE_DIR_BUNDLE_KEY in bundle: + out_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag + out_dir.mkdir(parents=True, exist_ok=True) + status_file = out_dir / "status.txt" + status_file.write_text( + f"status: {status_value}\n" + f"model_path: {raw_model_path}\n" + f"resolved_model_path: {resolved_model_path}\n", + encoding="utf-8", + ) + self.register_checkpoint(self.provides[0], status_file) + + bundle[f"{self.tag}.{self.provides[0]}"] = status_value diff --git a/TELF/pipeline/blocks/author_affiliation_tables.py b/TELF/pipeline/blocks/author_affiliation_tables.py new file mode 100644 index 00000000..6dce1f4d --- /dev/null +++ b/TELF/pipeline/blocks/author_affiliation_tables.py @@ -0,0 +1,141 @@ +# blocks/affiliations_and_authors_block.py +from __future__ import annotations + +from pathlib import Path +from typing import Dict, Sequence, Any, Optional, List, Tuple, Union + +import pandas as pd + +from .base_block import AnimalBlock +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY + +# your helpers live in blocks/block_helpers/ +from .block_helpers.affiliation_partition import generate_top_affiliations_with_country +from .block_helpers.author_partition import write_top_authors_by_cluster + + +class AffiliationsAndAuthorsBlock(AnimalBlock): + """ + Compute (affiliation, country, year) paper counts and top authors by cluster. + + ───────────────────────────────────────────────────────────── + always needs : ('df',) – accepts a CSV path OR a pandas.DataFrame + provides : ('affiliations_df', 'affiliations_csv', + 'authors_df', 'authors_csv') + tag : 'AffilsAndAuthors' (namespace for its outputs) + + Results are written under // . + Checkpoints persist the two CSV paths so the block can be skipped on re-run. + """ + + CANONICAL_NEEDS = ("df",) + + def __init__( + self, + *, + needs: Sequence[str] = CANONICAL_NEEDS, + provides: Sequence[str] = ("affiliations_df", "affiliations_csv", + "authors_df", "authors_csv"), + # Persist only the CSVs; DataFrames are rebuilt on load if needed. + checkpoint_keys: Sequence[str] = ("affiliations_csv", "authors_csv"), + conditional_needs: Sequence[tuple[str, Any]] = (), + tag: str = "AffilsAndAuthors", + # Defaults mirror your helper signatures + init_settings: Dict[str, Any] | None = None, + call_settings: Dict[str, Any] | None = None, + **kw: Any, + ) -> None: + + default_init: Dict[str, Any] = {} + + default_call: Dict[str, Any] = { + # generate_top_affiliations_with_country(...) + "min_total_papers": 20, + "country_filter": None, # exact match (including 'unknown') or None + "partition_by_year": False, + "per_year_output_dir": None, # if None and partition_by_year=True → use /by_year + + # write_top_authors_by_cluster(...) + "countries": None, # list[str] or None + "top_n": 10, + } + + super().__init__( + needs=needs, + provides=provides, + conditional_needs=list(conditional_needs or []), + checkpoint_keys=checkpoint_keys, + tag=tag, + init_settings=self._merge(default_init, init_settings), + call_settings=self._merge(default_call, call_settings), + **kw, + ) + + # ───────────────────────────────────────────────────────────── + # helpers + # ───────────────────────────────────────────────────────────── + def _ensure_input_csv(self, bundle: DataBundle) -> Path: + """ + Accepts either a DataFrame or a path in bundle['df']. + If a DataFrame, persist it to //input.csv and return that path. + """ + src = bundle[self.needs[0]] + save_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag + save_dir.mkdir(parents=True, exist_ok=True) + + if isinstance(src, pd.DataFrame): + inp = save_dir / "input.csv" + src.to_csv(inp, index=False, encoding="utf-8-sig") + return inp + + # let AnimalBlock’s path rewriter handle legacy numbered dirs + return Path(src) + + # ───────────────────────────────────────────────────────────── + # work + # ───────────────────────────────────────────────────────────── + def run(self, bundle: DataBundle) -> None: + df_path = self._ensure_input_csv(bundle) + + out_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag + out_dir.mkdir(parents=True, exist_ok=True) + + # === 1) Affiliations with country (and per-year optional) === + affils_csv = out_dir / "affiliations_top.csv" + per_year_dir = self.call_settings.get("per_year_output_dir") + if self.call_settings.get("partition_by_year") and not per_year_dir: + per_year_dir = out_dir / "by_year" + + generate_top_affiliations_with_country( + df_path=df_path, + affils_output_path=affils_csv, + min_total_papers=int(self.call_settings["min_total_papers"]), + country_filter=self.call_settings.get("country_filter"), + partition_by_year=bool(self.call_settings.get("partition_by_year")), + per_year_output_dir=per_year_dir, + ) + + # read back for the bundle + aff_df = pd.read_csv(affils_csv) if affils_csv.is_file() else pd.DataFrame( + columns=["affiliation_name", "country", "year", "paper_count"] + ) + + # register checkpoints / provide + self.register_checkpoint("affiliations_csv", affils_csv) + bundle[f"{self.tag}.affiliations_csv"] = str(affils_csv) + bundle[f"{self.tag}.affiliations_df"] = aff_df + + # === 2) Top authors by cluster ============================= + authors_csv = out_dir / "top_authors_by_cluster.csv" + # Note: helper param name is COUNTY_NAMES (kept as-is) + result_df = write_top_authors_by_cluster( + df_path=str(df_path), + output_path=str(authors_csv), + COUNTY_NAMES=self.call_settings.get("countries"), + top_n=int(self.call_settings.get("top_n", 10)), + ) + + # register checkpoints / provide + self.register_checkpoint("authors_csv", authors_csv) + bundle[f"{self.tag}.authors_csv"] = str(authors_csv) + bundle[f"{self.tag}.authors_df"] = result_df \ No newline at end of file diff --git a/TELF/pipeline/blocks/auto_bunny_simple_block.py b/TELF/pipeline/blocks/auto_bunny_simple_block.py old mode 100755 new mode 100644 diff --git a/TELF/pipeline/blocks/base_block.py b/TELF/pipeline/blocks/base_block.py index 81d2c96e..5549a92e 100644 --- a/TELF/pipeline/blocks/base_block.py +++ b/TELF/pipeline/blocks/base_block.py @@ -7,6 +7,8 @@ from PIL import Image import pickle from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY +import os, re, builtins +from contextlib import contextmanager Cond = Callable[[DataBundle, "AnimalBlock"], bool] @@ -52,6 +54,8 @@ def __init__( self._ckpt_keys = set(checkpoint_keys) if checkpoint_keys else set(self.provides) + + if hasattr(self.__class__, "CANONICAL_NEEDS"): self._canonical_needs = tuple(self.__class__.CANONICAL_NEEDS) else: @@ -114,6 +118,8 @@ def run(self, bundle: DataBundle) -> None: ... def __call__(self, bundle: DataBundle) -> DataBundle: + Path(bundle[SAVE_DIR_BUNDLE_KEY]).mkdir(parents=True, exist_ok=True) + # -------------------------------------------------------------- # 0) maybe load from checkpoint and bail early # -------------------------------------------------------------- @@ -124,7 +130,13 @@ def __call__(self, bundle: DataBundle) -> DataBundle: if exist: # hydrate only the keys we cached for k in self._ckpt_keys: - bundle[f"{self.tag}.{k}"] = self.load_path(ckpt[k]) + path = ckpt[k] + try: + bundle[f"{self.tag}.{k}"] = self.load_path(path) + except Exception: + # If reading (e.g., empty CSV) fails, just provide the path string. + # This still counts as "loaded", so run() is skipped. + bundle[f"{self.tag}.{k}"] = path # load any 'provides' from the class that is not saved to disk self._after_checkpoint_skip(bundle) @@ -150,7 +162,8 @@ def __call__(self, bundle: DataBundle) -> DataBundle: # 2) run block # -------------------------------------------------------------- self._pending_ckpt_map: Dict[str, str] = {} - self.run(bundle) + with self._io_path_rewriter(): + self.run(bundle) # -------------------------------------------------------------- # 3) verify outputs @@ -320,4 +333,111 @@ def copy(instance: Any, new_obj.needs = needs if provides: new_obj.provides = provides - return new_obj \ No newline at end of file + return new_obj + + + + def _base_tag(self, tag: str) -> str: + return re.sub(r'^\d+_', '', tag) + + @contextmanager + def _io_path_rewriter(self): + """ + While this block runs, rewrite any filesystem path segment matching + '/_/' (or '\\_\\') to the current self.tag, + for BOTH reads and writes. Also intercept directory creation to prevent + re-creating old numbered dirs. + """ + base_tag = self._base_tag(self.tag) + target_tag = self.tag + + # match start/middle/end occurrences for both / and \ + fwd_middle = re.compile(rf'(?<=/)\d+_{re.escape(base_tag)}(?=/)') + fwd_head = re.compile(rf'^\d+_{re.escape(base_tag)}(?=/)') + fwd_tail = re.compile(rf'(?<=/)\d+_{re.escape(base_tag)}$') + + bsl_middle = re.compile(rf'(?<=\\)\d+_{re.escape(base_tag)}(?=\\)') + bsl_head = re.compile(rf'^\d+_{re.escape(base_tag)}(?=\\)') + bsl_tail = re.compile(rf'(?<=\\)\d+_{re.escape(base_tag)}$') + + def _rewrite_seg(path_str: str) -> str: + s = path_str + s = fwd_middle.sub(target_tag, s) + s = fwd_head.sub(target_tag, s) + s = fwd_tail.sub(target_tag, s) + s = bsl_middle.sub(target_tag, s) + s = bsl_head.sub(target_tag, s) + s = bsl_tail.sub(target_tag, s) + return s + + # save originals + orig_open = builtins.open + orig_exists = os.path.exists + orig_isfile = os.path.isfile + orig_isdir = os.path.isdir + orig_mkdir = os.mkdir + orig_makedirs = os.makedirs + orig_path_mkdir = Path.mkdir + + def _open_patched(path, mode='r', *args, **kwargs): + try: + p = os.fspath(path) + except TypeError: + return orig_open(path, mode, *args, **kwargs) + p2 = _rewrite_seg(p) + # for reads: prefer rewritten, then fall back + if 'r' in mode and not any(ch in mode for ch in 'wax+'): + try: + return orig_open(p2, mode, *args, **kwargs) + except FileNotFoundError: + return orig_open(p, mode, *args, **kwargs) + # for writes or mixed modes: always use rewritten to avoid recreating old dirs + return orig_open(p2, mode, *args, **kwargs) + + def _wrap_exists(fn): + def _f(path): + try: + p = os.fspath(path) + except TypeError: + return fn(path) + return fn(_rewrite_seg(p)) + return _f + + def _mkdir_patched(path, mode=0o777): + p = _rewrite_seg(os.fspath(path)) + return orig_mkdir(p, mode) + + def _makedirs_patched(path, mode=0o777, exist_ok=False): + p = _rewrite_seg(os.fspath(path)) + return orig_makedirs(p, mode=mode, exist_ok=exist_ok) + + def _path_mkdir_patched(self_path: Path, mode=0o777, parents=False, exist_ok=False): + # emulate Path.mkdir but create at the rewritten location + p = _rewrite_seg(str(self_path)) + if parents: + return orig_makedirs(p, mode=mode, exist_ok=exist_ok) + else: + try: + return orig_mkdir(p, mode) + except FileExistsError: + if exist_ok: + return + raise + + try: + builtins.open = _open_patched + os.path.exists = _wrap_exists(orig_exists) + os.path.isfile = _wrap_exists(orig_isfile) + os.path.isdir = _wrap_exists(orig_isdir) + os.mkdir = _mkdir_patched + os.makedirs = _makedirs_patched + Path.mkdir = _path_mkdir_patched + yield + finally: + builtins.open = orig_open + os.path.exists = orig_exists + os.path.isfile = orig_isfile + os.path.isdir = orig_isdir + os.mkdir = orig_mkdir + os.makedirs = orig_makedirs + Path.mkdir = orig_path_mkdir diff --git a/TELF/pipeline/blocks/beaver_codependency_matrix_block.py b/TELF/pipeline/blocks/beaver_codependency_matrix_block.py index 67be7718..1261dc71 100644 --- a/TELF/pipeline/blocks/beaver_codependency_matrix_block.py +++ b/TELF/pipeline/blocks/beaver_codependency_matrix_block.py @@ -1,91 +1,230 @@ -# blocks/codependency_matrix_block.py from __future__ import annotations +# TELF/pipeline/blocks/beaver_codependency_matrix_block.py from pathlib import Path from typing import Any, Dict, Sequence, Tuple -import os, sparse -from ...pre_processing import Beaver -from ...helpers.file_system import load_file_as_dict +import numpy as np +import pickle +import scipy.sparse as sp +import sparse # pydata/sparse from .base_block import AnimalBlock from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY +# Ensure Beaver.coauthor_tensor is patched to a robust version on import +# from ...pre_processing import Beaver + +# TELF/pre_processing/Beaver/monkey_patch_coauthor_tensor.py +""" +Monkey-patch Beaver.coauthor_tensor to be robust to: +- n_jobs <= 0 (uses CPU count) +- empty / missing authors (writes an empty but valid tensor) +- missing 'year' column (uses 0) +Also saves authors/time index maps for downstream consumers. +""" + + +import os +import pickle +from typing import Dict, List, Tuple + +import numpy as np +import pandas as pd +import sparse # pydata/sparse + +from ...pre_processing import Beaver + + +def _safe_coauthor_tensor( + self: Beaver, + *, + dataset: pd.DataFrame, + target_columns: List[str], + split_authors_with: str = ";", + verbose: int = 0, + save_path: str | None = None, + n_nodes: int | None = None, + n_jobs: int = 1, + joblib_backend: str | None = None, + authors_idx_map: Dict[str, int] | None = None, + time_idx_map: Dict[int, int] | None = None, + return_object: bool = False, + output_mode: str | None = None, +): + auth_col, time_col = target_columns + df = dataset.copy() + + if time_col not in df.columns: + df[time_col] = 0 + + # Parse authors per doc + auth_series = ( + df[auth_col].fillna("") + .astype(str) + .str.split(split_authors_with) + .apply(lambda lst: [a.strip() for a in lst if a and a.strip()]) + ) + times = df[time_col].fillna(0).astype(int).tolist() + + # Build indices + if authors_idx_map is None: + unique_authors = sorted({a for lst in auth_series.tolist() for a in lst}) + a2i: Dict[str, int] = {a: i for i, a in enumerate(unique_authors)} + else: + a2i = dict(authors_idx_map) + + if time_idx_map is None: + unique_times = sorted(set(int(t) for t in times)) + if not unique_times: + unique_times = [0] + t2i: Dict[int, int] = {t: i for i, t in enumerate(unique_times)} + else: + t2i = dict(time_idx_map) + + A = len(a2i) + T = len(t2i) if t2i else 1 + + # Early exit: no authors → empty but valid 3D tensor (0 x 0 x max(1,T)) + if A == 0: + coo = sparse.COO(np.zeros((0, 0, T), dtype=np.float32)) + if save_path: + os.makedirs(save_path, exist_ok=True) + sparse.save_npz(os.path.join(save_path, "coauthor.npz"), coo) + with open(os.path.join(save_path, "authors_idx_map.p"), "wb") as f: + pickle.dump(a2i, f) + with open(os.path.join(save_path, "time_idx_map.p"), "wb") as f: + pickle.dump(t2i, f) + return coo if return_object else None + + # Build weighted undirected pairs per time + from collections import Counter + + weight = Counter() + for lst, t in zip(auth_series.tolist(), times): + idxs = [a2i[a] for a in lst if a in a2i] + if len(idxs) < 2: + continue + ti = t2i.get(int(t), next(iter(t2i.values())) if t2i else 0) + for i in range(len(idxs)): + for j in range(i + 1, len(idxs)): + u, v = idxs[i], idxs[j] + weight[(u, v, ti)] += 1 + weight[(v, u, ti)] += 1 # undirected + + if weight: + coords = np.array(list(zip(*weight.keys()))) + data = np.array(list(weight.values()), dtype=np.float32) + coo = sparse.COO(coords, data, shape=(A, A, T)) + else: + coo = sparse.COO(np.zeros((A, A, T), dtype=np.float32)) + + if save_path: + os.makedirs(save_path, exist_ok=True) + sparse.save_npz(os.path.join(save_path, "coauthor.npz"), coo) + with open(os.path.join(save_path, "authors_idx_map.p"), "wb") as f: + pickle.dump(a2i, f) + with open(os.path.join(save_path, "time_idx_map.p"), "wb") as f: + pickle.dump(t2i, f) + + return coo if return_object else None + + +# Apply the patch at import time +Beaver.coauthor_tensor = _safe_coauthor_tensor # type: ignore[misc] + class CodependencyMatrixBlock(AnimalBlock): """ - Build a 3-mode author–year tensor and flatten it to a co-authorship - matrix + node-ID map. + Build a (flattened) co-dependency matrix from a column of semicolon-separated ids + 'year'. - ───────────────────────────────────────────────────────────── - needs : ('df',) - provides : ('X', 'node_ids') - tag : 'CodeMatrix' + needs: ['df'] + provides: ['X', 'node_ids'] """ - CANONICAL_NEEDS = ('df', ) - # ------------------------------------------------------------------ # - # constructor # - # ------------------------------------------------------------------ # def __init__( self, *, - col: str = "slic_author_ids", - needs: Sequence[str] = CANONICAL_NEEDS, + col: str, + needs: Sequence[str] = ("df",), provides: Sequence[str] = ("X", "node_ids"), - conditional_needs: Sequence[Tuple[str, Any]] = (), # none for now - tag: str = "CodeMatrix", + tag: str = "BeaverCodependencyMatrix", + conditional_needs: Sequence[Tuple[str, Any]] = (), init_settings: Dict[str, Any] | None = None, call_settings: Dict[str, Any] | None = None, - verbose: bool = True, - **kwargs: Any, + **kw: Any, ) -> None: - - self.col = col # store the column name - - default_init = {} - default_call = { - "target_columns": [self.col, "year"], - "split_authors_with": ";", - "verbose": True, - "n_jobs": -1, - "authors_idx_map": {}, - "joblib_backend": "threading", - } - + self.col = col + default_init: Dict[str, Any] = {} + default_call: Dict[str, Any] = {"split_authors_with": ";", "n_jobs": 1} super().__init__( needs=needs, provides=provides, conditional_needs=conditional_needs, tag=tag, - init_settings=self._merge(default_init, init_settings), - call_settings=self._merge(default_call, call_settings), - verbose=verbose, - **kwargs, + init_settings={**default_init, **(init_settings or {})}, + call_settings={**default_call, **(call_settings or {})}, + **kw, ) - # ------------------------------------------------------------------ # - # work # - # ------------------------------------------------------------------ # def run(self, bundle: DataBundle) -> None: - # paths - out_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / "CodependencyMatrixBlock" / self.col - out_dir.mkdir(parents=True, exist_ok=True) + raw = bundle[self.needs[0]] + df = self.load_path(raw) if isinstance(raw, (str, Path)) else raw - # dataframe - df = bundle[self.needs[0]].copy() + # Ensure 'year' exists (Beaver expects it for the 3-mode tensor) + if "year" not in df.columns: + df = df.copy() + df["year"] = 0 - # build tensor with Beaver - beaver = Beaver(**self.init_settings) - cfg = dict(self.call_settings) - cfg.update({"dataset": df, "target_columns": [self.col, "year"], "save_path": out_dir}) + out_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag + out_dir.mkdir(parents=True, exist_ok=True) + # Run Beaver to materialize a 3D co-author tensor (A x A x T) + beaver = Beaver() + cfg = dict(self.call_settings) + cfg.update( + { + "dataset": df, + "target_columns": [self.col, "year"], + "save_path": str(out_dir), + } + ) beaver.coauthor_tensor(**cfg) - # load results - X = sparse.load_npz(out_dir / "coauthor.npz").sum(axis=2) # flatten 3-mode tensor - node_ids = load_file_as_dict(out_dir / "Authors.txt") - - # write back under this block’s namespace - bundle[f"{self.tag}.{self.provides[0]}"] = X - bundle[f"{self.tag}.{self.provides[1]}"] = node_ids + # Load the tensor and flatten across time + coo3: "sparse.COO" = sparse.load_npz(out_dir / "coauthor.npz") + # flatten T mode → 2D A x A + coo2 = coo3.sum(axis=2) + + # Convert to scipy.sparse (CSR) + if hasattr(coo2, "coords") and hasattr(coo2, "data"): + rows, cols = coo2.coords[0], coo2.coords[1] + X = sp.csr_matrix((coo2.data, (rows, cols)), shape=coo2.shape) + else: + # Fallback + X = sp.csr_matrix(np.asarray(coo2)) + + # Node id order from Beaver's authors_idx_map if present + a_map_path = out_dir / "authors_idx_map.p" + if a_map_path.exists(): + with open(a_map_path, "rb") as f: + a2i: Dict[str, int] = pickle.load(f) + node_ids = [None] * len(a2i) + for a, idx in a2i.items(): + node_ids[idx] = a + else: + # Fallback: derive from df (order might differ from Beaver) + ids = ( + df[self.col] + .dropna() + .astype(str) + .str.split(cfg.get("split_authors_with", ";")) + .explode() + .str.strip() + ) + ids = ids.loc[ids != ""].unique().tolist() + node_ids = sorted(set(ids)) + + # Publish outputs (top-level keys; WolfBlock expects this) + bundle[self.provides[0]] = X + bundle[self.provides[1]] = node_ids diff --git a/TELF/pipeline/blocks/block_helpers/KernelServer.py b/TELF/pipeline/blocks/block_helpers/KernelServer.py new file mode 100644 index 00000000..21e4368c --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/KernelServer.py @@ -0,0 +1,109 @@ +import os, signal, subprocess, pathlib, shlex, time + +class KernelTiedServer: + def __init__(self, cmd, *, cwd=None, env=None, log_dir=None): + """ + cmd: list or string (string is shell-split) + cwd: working dir + env: dict of env vars + log_dir: if set, stdout/stderr are appended to files here + """ + if isinstance(cmd, str): + cmd = shlex.split(cmd) + self.cmd = cmd + self.cwd = cwd + self.env = env + self.log_dir = pathlib.Path(log_dir).expanduser() if log_dir else None + self.proc = None + self.watchdog = None + self.stdout_f = None + self.stderr_f = None + + def start(self): + if self.proc and self.running: + raise RuntimeError("Already running") + if self.log_dir: + self.log_dir.mkdir(parents=True, exist_ok=True) + self.stdout_f = open(self.log_dir / "server.out", "ab", buffering=0) + self.stderr_f = open(self.log_dir / "server.err", "ab", buffering=0) + + # 1) start the server in its own process group + self.proc = subprocess.Popen( + self.cmd, + cwd=self.cwd, + env=self.env, + stdout=self.stdout_f or subprocess.DEVNULL, + stderr=self.stderr_f or subprocess.DEVNULL, + preexec_fn=os.setpgrp, # new PGID == proc.pid (POSIX) + ) + + # 2) watchdog: when kernel PID disappears, kill the whole PGID + kernel_pid = os.getpid() + pgid = self.proc.pid + script = f""" + while kill -0 {kernel_pid} 2>/dev/null; do sleep 2; done + kill -TERM -{pgid} 2>/dev/null + sleep 5 + kill -KILL -{pgid} 2>/dev/null + """ + self.watchdog = subprocess.Popen( + ["bash", "-c", script], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ) + return self + + @property + def running(self): + return self.proc is not None and self.proc.poll() is None + + def status(self): + if not self.proc: + return "not started" + rc = self.proc.poll() + return f"running (pid {self.proc.pid}, pgid {self.proc.pid})" if rc is None else f"exited rc={rc}" + + def stop(self, sig=signal.SIGTERM, hard_after=5): + if not self.proc: + return + if self.running: + try: + os.killpg(self.proc.pid, sig) + except ProcessLookupError: + pass + # optional hard kill after grace + t0 = time.time() + while self.running and time.time() - t0 < hard_after: + time.sleep(0.2) + if self.running: + try: + os.killpg(self.proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + # stop watchdog + if self.watchdog and self.watchdog.poll() is None: + try: + self.watchdog.terminate() + except Exception: + pass + # close logs + for f in (self.stdout_f, self.stderr_f): + try: + f and f.close() + except Exception: + pass + + def tail(self, n=50, which="out"): + if not self.log_dir: + print("(no logs: set log_dir=...)") + return + path = self.log_dir / ( "server.out" if which=="out" else "server.err" ) + try: + with open(path, "rb") as f: + print(b"".join(f.readlines()[-n:]).decode(errors="replace")) + except FileNotFoundError: + print("(no log yet)") + +# --- Example --- +# srv = KernelTiedServer(["zsh", "../../Lynx/start_lynx.sh"], log_dir="~/.logs/lynx").start() +# print(srv.status()); srv.tail(100) # view last 100 lines of stdout +# srv.stop() # stop manually diff --git a/TELF/pipeline/blocks/block_helpers/__init__.py b/TELF/pipeline/blocks/block_helpers/__init__.py new file mode 100644 index 00000000..c2e9f682 --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/__init__.py @@ -0,0 +1 @@ +from .KernelServer import KernelTiedServer \ No newline at end of file diff --git a/TELF/pipeline/blocks/block_helpers/affiliation_partition.py b/TELF/pipeline/blocks/block_helpers/affiliation_partition.py new file mode 100644 index 00000000..2fab77a3 --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/affiliation_partition.py @@ -0,0 +1,256 @@ +# TELF/pipeline/blocks/block_helpers/affiliation_partition.py +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +import ast +import json +import numpy as np +import pandas as pd + +UNKNOWN_COUNTRY = "Unknown" + + +def _parse_affiliations_field(val: Any) -> Dict[str, Dict[str, Any]]: + """ + Normalize an affiliations field into a dict-of-dicts keyed by affiliation id. + + Accepts: + - dict (already keyed by id) with nested dicts + - list[dict] (each item is an affiliation) + - string (Python-literal or JSON) representing either of the above + - else -> empty dict + + Ensures each nested dict has keys: + - 'name' (str) + - 'country' (str, default 'Unknown') + """ + # 1) Convert to Python object + obj = None + if isinstance(val, (dict, list)): + obj = val + elif isinstance(val, str): + s = val.strip() + if s: + # Try Python literal first (TELF often saves repr strings) + for parser in (ast.literal_eval, json.loads): + try: + parsed = parser(s) + if isinstance(parsed, (dict, list)): + obj = parsed + break + except Exception: + pass + if obj is None: + obj = {} + else: + obj = {} + + # 2) Canonicalize to dict-of-dicts keyed by id (string keys) + out: Dict[str, Dict[str, Any]] = {} + if isinstance(obj, dict): + for k, v in obj.items(): + if not isinstance(v, dict): + continue + name = v.get("name") + country = v.get("country", UNKNOWN_COUNTRY) + if name is None or (isinstance(name, str) and name.strip() == ""): + # try fallbacks (some loaders store name under different keys) + name = v.get("affiliation_name") or v.get("org") or v.get("institution") or None + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + out[str(k)] = {**v, "name": name, "country": country} + return out + + if isinstance(obj, list): + for i, item in enumerate(obj): + if not isinstance(item, dict): + continue + key = str(item.get("id", item.get("affiliation_id", i))) + name = item.get("name") or item.get("affiliation_name") or item.get("org") or item.get("institution") + country = item.get("country", UNKNOWN_COUNTRY) + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + out[key] = {**item, "name": name, "country": country} + return out + + return {} + + +def _pairs_from_affiliations(val: Any) -> List[Tuple[Optional[str], str]]: + """ + Produce a list of (affiliation_name, country) pairs from a raw affiliations field. + Ensures each pair has exactly two elements; missing pieces are filled with defaults. + """ + norm = _parse_affiliations_field(val) + pairs: List[Tuple[Optional[str], str]] = [] + for _, rec in norm.items(): + name = rec.get("name") + if isinstance(name, str): + name = name.strip() or None + country = rec.get("country", UNKNOWN_COUNTRY) + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + pairs.append((name, country)) + return pairs + + +def _resolve_paper_id_column(df: pd.DataFrame) -> str: + """ + Choose a robust paper-id column for grouping: + priority: 'eid' -> 's2id' -> 'doi' -> synthetic 'paper_id' + Returns the *name* of the column (and creates synthetic if needed). + """ + for cand in ("eid", "s2id", "doi"): + if cand in df.columns: + return cand + df = df.reset_index(drop=True) + df["paper_id"] = df.index.astype(str) + return "paper_id" + + +def generate_top_affiliations_with_country( + df_path: Union[str, Path], + affils_output_path: Union[str, Path], + min_total_papers: int = 1, + country_filter: Optional[str] = None, + partition_by_year: bool = False, + per_year_output_dir: Optional[Union[str, Path]] = None, +) -> None: + """ + Read the (already processed) pipeline CSV and emit a table of top affiliations + with countries, optionally partitioned by year. + + Output schema (affils_output_path): + ['affiliation_name', 'country', 'year', 'paper_count'] + + Parameters + ---------- + df_path : path to the dataframe (CSV) produced upstream + affils_output_path : where to write the aggregated CSV + min_total_papers : keep only affiliations with >= this many papers overall + country_filter : if given, keep only rows matching this country (case-sensitive) + partition_by_year : if True, also write per-year CSVs under per_year_output_dir + per_year_output_dir : base directory for per-year outputs; if None and + partition_by_year is True, defaults to affils_output_path.parent / 'by_year' + """ + df_path = Path(df_path) + out_path = Path(affils_output_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + + if not df_path.is_file(): + # nothing to do; write an empty file with the expected header + empty = pd.DataFrame(columns=["affiliation_name", "country", "year", "paper_count"]) + empty.to_csv(out_path, index=False) + return + + df = pd.read_csv(df_path) + + # pick affiliations column (prefer SLIC) + aff_col = "slic_affiliations" if "slic_affiliations" in df.columns else ( + "affiliations" if "affiliations" in df.columns else None + ) + if aff_col is None: + # no affiliations at all -> write empty + empty = pd.DataFrame(columns=["affiliation_name", "country", "year", "paper_count"]) + empty.to_csv(out_path, index=False) + return + + # ensure year column + if "year" not in df.columns: + df["year"] = 0 + else: + df["year"] = pd.to_numeric(df["year"], errors="coerce").fillna(0).astype(int) + + # robust paper id + pid_col = _resolve_paper_id_column(df) + if pid_col not in df.columns: + # _resolve_paper_id_column may have added a synthetic col; ensure present + df = df.reset_index(drop=True) + df["paper_id"] = df.index.astype(str) + pid_col = "paper_id" + + # Build normalized pairs per row and explode + df = df[[pid_col, "year", aff_col]].copy() + df["affil_pairs"] = df[aff_col].apply(_pairs_from_affiliations) + + # explode to one row per (paper, affiliation) + exploded = df.explode("affil_pairs") + + # Normalize pairs so every row is a *2-tuple* (name, country) + def _safe_pair(x: Any) -> Tuple[Optional[str], str]: + if isinstance(x, (list, tuple)): + if len(x) >= 2: + name, country = x[0], x[1] + elif len(x) == 1: + name, country = x[0], UNKNOWN_COUNTRY + else: + name, country = None, UNKNOWN_COUNTRY + elif isinstance(x, dict): + name = x.get("name") + country = x.get("country", UNKNOWN_COUNTRY) + else: + name, country = None, UNKNOWN_COUNTRY + + if isinstance(name, str): + name = name.strip() or None + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + return (name, country) + + exploded["affil_pairs"] = exploded["affil_pairs"].apply(_safe_pair) + # Now guaranteed to be 2 columns + exploded[["affiliation_name", "country"]] = pd.DataFrame( + exploded["affil_pairs"].tolist(), index=exploded.index + ) + + # Drop rows where affiliation name is missing after normalization + exploded = exploded.dropna(subset=["affiliation_name"]).copy() + + # Optional country filter + if country_filter: + exploded = exploded.loc[exploded["country"] == str(country_filter)].copy() + + if exploded.empty: + out = pd.DataFrame(columns=["affiliation_name", "country", "year", "paper_count"]) + out.to_csv(out_path, index=False) + return + + # Unique by (paper, affiliation) to avoid double-counting the same affiliation within a paper + exploded = exploded[[pid_col, "year", "affiliation_name", "country"]].drop_duplicates() + + # Compute total papers per affiliation across all years (for thresholding) + totals = ( + exploded.groupby(["affiliation_name", "country"])[pid_col] + .nunique() + .reset_index(name="paper_count_total") + ) + + # Keep only affiliations that meet the threshold + keep = totals.loc[totals["paper_count_total"] >= int(min_total_papers), ["affiliation_name", "country"]] + if keep.empty: + out = pd.DataFrame(columns=["affiliation_name", "country", "year", "paper_count"]) + out.to_csv(out_path, index=False) + return + + # Join to filter + exploded = exploded.merge(keep, on=["affiliation_name", "country"], how="inner") + + # Aggregate by year + by_year = ( + exploded.groupby(["affiliation_name", "country", "year"])[pid_col] + .nunique() + .reset_index(name="paper_count") + .sort_values(["paper_count", "year"], ascending=[False, True]) + .reset_index(drop=True) + ) + + by_year.to_csv(out_path, index=False) + + # Optionally write per-year partitions + if partition_by_year: + base = Path(per_year_output_dir) if per_year_output_dir else out_path.parent / "by_year" + base.mkdir(parents=True, exist_ok=True) + for yr, sub in by_year.groupby("year"): + sub_path = base / f"affiliations_{yr}.csv" + sub.sort_values("paper_count", ascending=False).to_csv(sub_path, index=False) diff --git a/TELF/pipeline/blocks/block_helpers/author_partition.py b/TELF/pipeline/blocks/block_helpers/author_partition.py new file mode 100644 index 00000000..c90ce5d9 --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/author_partition.py @@ -0,0 +1,302 @@ +# TELF/pipeline/blocks/block_helpers/author_partition.py +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +import ast +import json +import numpy as np +import pandas as pd + +UNKNOWN_COUNTRY = "Unknown" + + +# ----------------------------- Parsers --------------------------------- + +def _parse_literal_or_json(text: str): + """Try ast.literal_eval then JSON; return Python object or None.""" + if not isinstance(text, str) or not text.strip(): + return None + s = text.strip() + for parser in (ast.literal_eval, json.loads): + try: + return parser(s) + except Exception: + pass + return None + + +def _to_list_any(x: Any) -> List[Any]: + """Normalize a value to a Python list (best‑effort).""" + if x is None or (isinstance(x, float) and pd.isna(x)): + return [] + if isinstance(x, (list, tuple)): + return list(x) + if isinstance(x, str): + parsed = _parse_literal_or_json(x) + if isinstance(parsed, (list, tuple)): + return list(parsed) + # fall back to delimiter split + sep = ";" if ";" in x else "," + return [t.strip() for t in x.split(sep) if t.strip()] + return [x] + + +def _parse_authors_field(row: pd.Series) -> List[str]: + """ + Return a list of author IDs as strings. + Tries SLIC → generic → S2 IDs; falls back to empty. + """ + for cand in ("slic_author_ids", "author_ids", "s2_author_ids"): + if cand in row and pd.notna(row[cand]): + vals = _to_list_any(row[cand]) + return [str(v).strip() for v in vals if str(v).strip() != ""] + return [] + + +def _author_id_to_name_map(row: pd.Series) -> Dict[str, str]: + """ + If both author IDs and names exist with the same length, return an id→name map. + Otherwise return {}. + """ + id_cols = [c for c in ("slic_author_ids", "author_ids", "s2_author_ids") if c in row.index] + name_cols = [c for c in ("slic_authors", "authors") if c in row.index] + for ic in id_cols: + for nc in name_cols: + ids = _to_list_any(row.get(ic)) + names = _to_list_any(row.get(nc)) + if len(ids) == len(names) and len(ids) > 0: + out = {} + for i, n in zip(ids, names): + sid = str(i).strip() + name = str(n).strip() + if sid: + out[sid] = name + if out: + return out + return {} + + +def _parse_affiliations_field(val: Any) -> Dict[str, Dict[str, Any]]: + """ + Normalize an affiliations field into a dict-of-dicts keyed by affiliation id (string). + + Accepts dict, list[dict], or str (Python‑literal or JSON). Returns a dict where each + value has at least: + { "name": str|None, "country": str, "authors": List[str] } + """ + # 1) Convert to Python object + if isinstance(val, (dict, list)): + obj = val + elif isinstance(val, str): + obj = _parse_literal_or_json(val) or {} + else: + obj = {} + + out: Dict[str, Dict[str, Any]] = {} + + if isinstance(obj, dict): + items = obj.items() + elif isinstance(obj, list): + # fabricate keys if missing + items = [(str(i.get("id", i.get("affiliation_id", idx))), i) + for idx, i in enumerate(obj) if isinstance(i, dict)] + else: + items = [] + + for k, v in items: + if not isinstance(v, dict): + continue + name = v.get("name") or v.get("affiliation_name") or v.get("org") or v.get("institution") + if isinstance(name, str) and not name.strip(): + name = None + country = v.get("country", UNKNOWN_COUNTRY) + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + + # authors may be under various keys + auths = v.get("authors", v.get("author_ids", [])) + auth_list = [str(a).strip() for a in _to_list_any(auths) if str(a).strip() != ""] + out[str(k)] = {"name": name, "country": country, "authors": auth_list, **v} + + return out + + +def _choose_paper_id_column(df: pd.DataFrame) -> str: + """Pick a stable paper-id column for de-duplication and counting.""" + for cand in ("eid", "s2id", "doi"): + if cand in df.columns: + return cand + # synthesize one + if "paper_id" not in df.columns: + df["paper_id"] = np.arange(len(df)).astype(str) + return "paper_id" + + +def _most_common_non_null(series: pd.Series) -> Any: + """Return the most frequent non-null value in a Series, or np.nan.""" + s = series.dropna() + if s.empty: + return np.nan + return s.value_counts().idxmax() + + +# ------------------------- Public API ---------------------------------- + +def write_top_authors_by_cluster( + df_path: Union[str, Path], + output_path: Union[str, Path], + COUNTY_NAMES: Optional[Iterable[str]] = None, + top_n: int = 10, + debug: bool = False, +) -> pd.DataFrame: + """ + Build a 'top authors by cluster' table. + + Output CSV schema: + ['cluster', 'author_id', 'author', 'affiliation_name', 'country', + 'paper_count', 'num_citations'] + + Notes + ----- + * COUNTY_NAMES is preserved for backward compatibility in the caller. If given, + rows are filtered to only those countries. + * Per-paper credit: each author gets 1 paper for that row (unique paper_id), + and the row's num_citations are summed across their papers. + """ + df_path = Path(df_path) + out_path = Path(output_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + + # If no input, write empty CSV with expected header. + if not df_path.is_file(): + empty = pd.DataFrame( + columns=["cluster", "author_id", "author", "affiliation_name", "country", + "paper_count", "num_citations"] + ) + empty.to_csv(out_path, index=False) + return empty + + df = pd.read_csv(df_path) + + # Ensure basic columns + if "cluster" not in df.columns: + df["cluster"] = 0 + if "num_citations" not in df.columns: + df["num_citations"] = 0 + + # Choose paper id column robustly + pid_col = _choose_paper_id_column(df) + + # Prefer SLIC affiliations, fall back + aff_col = "slic_affiliations" if "slic_affiliations" in df.columns else ( + "affiliations" if "affiliations" in df.columns else None + ) + + # Build records + records: List[Dict[str, Any]] = [] + + for _, row in df.iterrows(): + cluster = int(row.get("cluster", 0)) if pd.notna(row.get("cluster", np.nan)) else 0 + paper_id = str(row.get(pid_col, "")) + citations = row.get("num_citations", 0) + try: + citations = float(citations) + except Exception: + citations = 0.0 + + author_ids = _parse_authors_field(row) + id2name = _author_id_to_name_map(row) + + # Parse affiliations + create author→(name,country) lookup + aff_map: Dict[str, Dict[str, Any]] = {} + if aff_col is not None and aff_col in row and pd.notna(row[aff_col]): + aff_map = _parse_affiliations_field(row[aff_col]) + + # Build reverse index: author_id -> list of (aff_name, country) + author_to_aff: Dict[str, List[Tuple[Optional[str], str]]] = {} + for _, info in aff_map.items(): + aff_name = info.get("name") + country = info.get("country", UNKNOWN_COUNTRY) + if not isinstance(country, str) or not country.strip(): + country = UNKNOWN_COUNTRY + for aid in info.get("authors", []): + author_to_aff.setdefault(str(aid), []).append((aff_name, country)) + + # Create a record per author + for aid in author_ids: + aff_pairs = author_to_aff.get(aid, []) + if aff_pairs: + # pick most common (name,country) for this paper row + names = pd.Series([a for a, _ in aff_pairs], dtype="object") + cntrs = pd.Series([c for _, c in aff_pairs], dtype="object") + aff_name = _most_common_non_null(names) + country = _most_common_non_null(cntrs) + else: + aff_name = np.nan + country = UNKNOWN_COUNTRY + + records.append( + dict( + cluster=cluster, + paper_id=paper_id, + author_id=str(aid), + author=id2name.get(str(aid), np.nan), + affiliation_name=aff_name if (isinstance(aff_name, str) and aff_name.strip()) else np.nan, + country=country if (isinstance(country, str) and country.strip()) else UNKNOWN_COUNTRY, + num_citations=citations, + ) + ) + + if not records: + empty = pd.DataFrame( + columns=["cluster", "author_id", "author", "affiliation_name", "country", + "paper_count", "num_citations"] + ) + empty.to_csv(out_path, index=False) + return empty + + rec_df = pd.DataFrame.from_records(records) + + # Optionally filter by country list (param name preserved as in caller) + if COUNTY_NAMES: + counties = {str(c).strip() for c in COUNTY_NAMES if str(c).strip()} + if counties: + rec_df = rec_df[rec_df["country"].isin(counties)].copy() + + if rec_df.empty: + empty = pd.DataFrame( + columns=["cluster", "author_id", "author", "affiliation_name", "country", + "paper_count", "num_citations"] + ) + empty.to_csv(out_path, index=False) + return empty + + # Aggregate: unique paper count per (cluster, author_id) and sum citations + agg = ( + rec_df.groupby(["cluster", "author_id"], dropna=False) + .agg( + paper_count=("paper_id", "nunique"), + num_citations=("num_citations", "sum"), + # pick most common non-null strings + author=("author", _most_common_non_null), + affiliation_name=("affiliation_name", _most_common_non_null), + country=("country", _most_common_non_null), + ) + .reset_index() + ) + + # Rank within each cluster + agg["num_citations"] = pd.to_numeric(agg["num_citations"], errors="coerce").fillna(0.0) + agg["paper_count"] = pd.to_numeric(agg["paper_count"], errors="coerce").fillna(0).astype(int) + + agg = agg.sort_values(["cluster", "num_citations", "paper_count"], ascending=[True, False, False]) + + # Keep top_n per cluster + top = agg.groupby("cluster", group_keys=False).head(int(top_n)).reset_index(drop=True) + + # Write and return + top.to_csv(out_path, index=False) + if debug: + print(f"[author_partition] Wrote top authors by cluster → {out_path} (rows={len(top)})") + return top diff --git a/TELF/pipeline/blocks/block_helpers/hnmfk_paths.py b/TELF/pipeline/blocks/block_helpers/hnmfk_paths.py new file mode 100644 index 00000000..98e627af --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/hnmfk_paths.py @@ -0,0 +1,51 @@ +# post_processing/Peacock/node_sources.py +from __future__ import annotations +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Iterator, Optional + + +from ....factorization import HNMFk # adjust relative path if needed + + +@dataclass(frozen=True) +class Node: + dir: Path + csv: Path + + +class NodeSource: + def iter_nodes(self) -> Iterator[Node]: + raise NotImplementedError + + +class HNMFkNodeSource(NodeSource): + def __init__(self, experiment_path: Path, *, only_existing_csv: bool = True) -> None: + self.experiment_path = Path(experiment_path).expanduser().resolve() + self.only_existing_csv = only_existing_csv + + def iter_nodes(self) -> Iterator[Node]: + model = HNMFk(experiment_name=str(self.experiment_path)) + model.load_model() + seen = set() + for node in model.traverse_nodes(): + node_dir = Path(node["node_save_path"]).resolve().parent + if node_dir in seen: + continue + seen.add(node_dir) + # prefer latest cluster file in that directory + csvs = sorted(node_dir.glob("cluster_for_k=*.csv"), key=lambda p: p.stat().st_mtime) + if csvs: + yield Node(dir=node_dir, csv=csvs[-1]) + elif not self.only_existing_csv: + # allow “expected” path even if it doesn’t exist yet + yield Node(dir=node_dir, csv=node_dir / "cluster_for_k=UNKNOWN.csv") + + +class GlobNodeSource(NodeSource): + def __init__(self, root: Path) -> None: + self.root = Path(root).expanduser().resolve() + + def iter_nodes(self) -> Iterator[Node]: + for csv in sorted(self.root.rglob("cluster_for_k=*.csv")): + yield Node(dir=csv.parent, csv=csv) diff --git a/TELF/pipeline/blocks/block_helpers/peacock_renderer.py b/TELF/pipeline/blocks/block_helpers/peacock_renderer.py new file mode 100644 index 00000000..6726e6db --- /dev/null +++ b/TELF/pipeline/blocks/block_helpers/peacock_renderer.py @@ -0,0 +1,238 @@ +# post_processing/Peacock/peacock_renderer.py +from __future__ import annotations +from pathlib import Path +from typing import Dict, Optional, Sequence +import ast +import json +import numpy as np +import pandas as pd + +from ....post_processing.Peacock.Utility import aggregate_ostats +from ....post_processing.Peacock.Plot import plot_heatmap, plot_bar, plot_scatter + +plot_hist = plot_bar +plot_scatter3D = plot_scatter + + +def _normalize_aff_to_py_literal(v): + """ + Accept dict/list/JSON/string; return a *Python-literal* string (repr), + guaranteeing per-affiliation 'authors' (list) & 'country' (str). + """ + if isinstance(v, (dict, list)): + obj = v + elif isinstance(v, str): + s = v.strip() + if not s: + obj = [] + else: + try: + obj = ast.literal_eval(s) + except Exception: + try: + obj = json.loads(s) + except Exception: + obj = [] + else: + obj = [] + + if isinstance(obj, list): + out = {} + for i, item in enumerate(obj): + if not isinstance(item, dict): + continue + key = str(item.get("id", item.get("affiliation_id", i))) + out[key] = { + **item, + "authors": item.get("authors", item.get("author_ids", [])) or [], + "country": item.get("country", "Unknown"), + } + obj = out + elif isinstance(obj, dict): + out = {} + for k, val in obj.items(): + if not isinstance(val, dict): + continue + val.setdefault("authors", val.get("author_ids", [])) + val.setdefault("country", "Unknown") + out[str(k)] = val + obj = out + else: + obj = {} + + return repr(obj) + + +class PeacockRenderer: + def __init__( + self, + *, + hist_stats: Sequence[str] = ("paper_count", "num_citations"), + hist_ylabels: Optional[Dict[str, str]] = None, + col_names: Optional[Dict[str, str]] = None, + affiliation_palette: Optional[Dict[str, str]] = None, + country: Optional[str] = None, + cluster_col: Optional[str] = "cluster", # NEW + ) -> None: + self.hist_stats = tuple(hist_stats) + self.hist_ylabels = hist_ylabels or { + "paper_count": "Number of Papers", + "num_citations": "Number of Citations", + "attribution_percentage": "Attribution Percentage", + } + self.col_names = col_names or { + "id": "eid", + "authors": "slic_authors", + "author_ids": "slic_author_ids", + "affiliations": "slic_affiliations", + "funding": "funding", + "citations": "num_citations", + "references": "references", + } + self.affiliation_palette = affiliation_palette or {} + self.country = country + self.cluster_col = cluster_col + + def _png_or_html(self, make_plot_func, stem: Path, *args, **kwargs): + png_path = stem.with_suffix(".png") + html_path = stem.with_suffix(".html") + try: + fig = make_plot_func(*args, interactive=True, fname=None, **kwargs) + fig.write_html(str(html_path), include_plotlyjs="cdn") + except Exception: + print("Exception making interactive plots in peacock") + make_plot_func(*args, interactive=False, fname=str(png_path), **kwargs) + + def render(self, df: pd.DataFrame, out_dir: Path) -> None: + # run overall + self._render_core(df, out_dir) + + # run per-cluster + if self.cluster_col and self.cluster_col in df.columns: + cluster_root = out_dir / "clusters" + for cid, df_c in df.groupby(self.cluster_col, dropna=False): + safe = "nan" if pd.isna(cid) else str(cid).replace("/", "_") + self._render_core(df_c, cluster_root / safe) + + def _render_core(self, df: pd.DataFrame, out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + df = df.copy() + + aff_col = self.col_names["affiliations"] + aut_col = self.col_names["authors"] + aid_col = self.col_names["author_ids"] + + def _to_list_any(x): + if x is None or (isinstance(x, float) and pd.isna(x)): + return [] + if isinstance(x, (list, tuple)): + return list(x) + if isinstance(x, str): + s = x.strip() + if not s: + return [] + for parser in (ast.literal_eval, json.loads): + try: + v = parser(s) + if isinstance(v, list): + return v + except Exception: + pass + sep = ";" if ";" in s else "," + return [t.strip() for t in s.split(sep) if t.strip()] + return [str(x)] + + def _to_sc_str_preserve_nan(x): + lst = _to_list_any(x) + vals = [str(v).strip() for v in lst if str(v).strip() != ""] + return np.nan if not vals else ";".join(vals) + + df[aut_col] = df[aut_col].apply(_to_sc_str_preserve_nan) + df[aid_col] = df[aid_col].apply(_to_sc_str_preserve_nan) + + subset = [self.col_names["id"], aut_col, aid_col, aff_col] + df = df.dropna(subset=subset) + + if "year" in df.columns: + df["year"] = pd.to_numeric(df["year"], errors="coerce") + df = df.dropna(subset=["year"]) + df["year"] = df["year"].astype(int) + else: + df["year"] = 0 + + if df.empty: + (out_dir / "top_authors.csv").write_text("") + (out_dir / "top_affiliations.csv").write_text("") + return + + df[aff_col] = df[aff_col].apply(_normalize_aff_to_py_literal) + + filters = {"country": self.country} if self.country else None + + def _safe_pivot_table(data, index, columns, values): + if data is None or len(data) == 0: + return pd.DataFrame() + return data.pivot_table(index=index, columns=columns, values=values, aggfunc="sum", fill_value=0) + + author_stats = aggregate_ostats(df, key="author_id", top_n=100, col_names=self.col_names, filters=filters, by_year=False) + affiliation_stats = aggregate_ostats(df, key="affiliation_id", top_n=100, col_names=self.col_names, filters=filters, by_year=False) + author_stats.to_csv(out_dir / "top_authors.csv", index=False) + affiliation_stats.to_csv(out_dir / "top_affiliations.csv", index=False) + + auth_args = dict(key="author_id", top_n=10, sort_by="num_citations", col_names=self.col_names, by_year=True, filters=filters) + aff_args = dict(key="affiliation_id", top_n=10, sort_by="num_citations", col_names=self.col_names, by_year=True, filters=filters) + + auth_heat = aggregate_ostats(df, **auth_args) + aff_heat = aggregate_ostats(df, **aff_args) + + pivot_c = _safe_pivot_table(auth_heat, index="year", columns="author", values="num_citations") + pivot_p = _safe_pivot_table(auth_heat, index="year", columns="author", values="paper_count") + pivot_c2 = _safe_pivot_table(aff_heat, index="year", columns="affiliation", values="num_citations") + pivot_p2 = _safe_pivot_table(aff_heat, index="year", columns="affiliation", values="paper_count") + + if not pivot_c.empty: + self._png_or_html(plot_heatmap, out_dir / "author_heatmap_citations", + pivot_c, cmap="jet", interpolation="gaussian", + title="Author Citations by Year", xlabel="Author", ylabel="Year") + if not pivot_p.empty: + self._png_or_html(plot_heatmap, out_dir / "author_heatmap_papers", + pivot_p, cmap="jet", interpolation="gaussian", + title="Author Papers by Year", xlabel="Author", ylabel="Year") + if not pivot_c2.empty: + self._png_or_html(plot_heatmap, out_dir / "affiliation_heatmap_citations", + pivot_c2, cmap="jet", interpolation="gaussian", + title="Affiliation Citations by Year", xlabel="Affiliation", ylabel="Year") + if not pivot_p2.empty: + self._png_or_html(plot_heatmap, out_dir / "affiliation_heatmap_papers", + pivot_p2, cmap="jet", interpolation="gaussian", + title="Affiliation Papers by Year", xlabel="Affiliation", ylabel="Year") + + auth_hist = aggregate_ostats(df, **{**auth_args, "by_year": False}) + if not auth_hist.empty: + self._png_or_html(plot_hist, out_dir / "author_hist", + auth_hist, x="author", ys=list(self.hist_stats), + title="Author Statistics Histogram", xlabel="Author", + ylabel=self.hist_ylabels[self.hist_stats[0]]) + + aff_hist = aggregate_ostats(df, **{**aff_args, "by_year": False}) + if not aff_hist.empty: + self._png_or_html(plot_hist, out_dir / "affiliation_hist", + aff_hist, x="affiliation", ys=list(self.hist_stats), + title="Affiliation Statistics Histogram", xlabel="Affiliation", + ylabel=self.hist_ylabels[self.hist_stats[0]]) + + self._png_or_html(plot_scatter3D, out_dir / "author_scatter", + df, x="paper_count", y="attribution_percentage", z="num_citations", + agg_func=aggregate_ostats, agg_kwargs=auth_args, + log_z=True, hue="affiliation", labels="author", + title="Author Stats Scatter3D", xlabel="Paper Count", + ylabel="Attribution Percentage", zlabel="Num. Citations", + base_palette=self.affiliation_palette) + + self._png_or_html(plot_scatter3D, out_dir / "affiliation_scatter", + df, x="paper_count", y="attribution_percentage", z="num_citations", + agg_func=aggregate_ostats, agg_kwargs=aff_args, + log_z=True, hue="country", labels="affiliation", + title="Affiliation Stats Scatter3D", xlabel="Paper Count", + ylabel="Attribution Percentage", zlabel="Num. Citations", + base_palette=self.affiliation_palette) diff --git a/TELF/pipeline/blocks/collect_hnmfk_leaf_block.py b/TELF/pipeline/blocks/collect_hnmfk_leaf_block.py new file mode 100644 index 00000000..70e0712c --- /dev/null +++ b/TELF/pipeline/blocks/collect_hnmfk_leaf_block.py @@ -0,0 +1,447 @@ +# TELF/pipeline/blocks/leaf_data_labels_block.py +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Sequence, Tuple, Optional, List +import os +import re +import pickle +import pandas as pd +import numpy as np +from datetime import datetime + +from .base_block import AnimalBlock +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY + + +class CollectHNMFkLeafBlock(AnimalBlock): + """ + Build two artifacts from a completed HNMFk run: + + • LEAF_DATA.csv + ALL columns from the input df for documents in each LEAF node, + plus: cluster (int), Graph_Name (e.g. 'depth_2_Parent_1_0') + + • LEAF_LABELS.csv + Columns: Graph_Name, label, words + (words are comma-joined from top_words.csv for each cluster) + + • summary.txt + Total documents, number of leaf clusters, and per-cluster counts. + + Resolution order for HNMFk experiment directory: + 1) From bundle key (default: 'SemanticHNMFk.model_path', overridable) + 2) call_settings['hnmfk_dir'] + 3) Auto-discover under SAVE_DIR_BUNDLE_KEY using tag (default: 'SemanticHNMFk') + + Optional call_settings: + - hnmfk_dir: explicit path to the experiment directory + - hnmfk_tag: tag name to look for (default: 'SemanticHNMFk') + - hnmfk_bundle_key: bundle key for model path (default: f'{hnmfk_tag}.model_path') + + Provides: + - 'leaf_data_csv' → Path to LEAF_DATA.csv + - 'leaf_labels_csv' → Path to LEAF_LABELS.csv + (summary.txt is written alongside these outputs) + """ + + CANONICAL_NEEDS: Tuple[str, ...] = ("df", SAVE_DIR_BUNDLE_KEY) + + def __init__( + self, + *, + needs: Sequence[str] = CANONICAL_NEEDS, + provides: Sequence[str] = ("df", "leaf_labels_csv"), + tag: str = "LeafDataLabels", + init_settings: Optional[Dict[str, Any]] = None, + call_settings: Optional[Dict[str, Any]] = None, + verbose: bool = True, + **kw: Any, + ) -> None: + super().__init__( + needs=needs, + provides=provides, + tag=tag, + init_settings=init_settings or {}, + call_settings=call_settings or {}, + verbose=verbose, + checkpoint=True, + load_checkpoint=False, # force a real run the first time + **kw, + ) + + # ───────────────────────── helpers ───────────────────────── + + @staticmethod + def _to_list(x, *, as_int=False, as_str=False) -> List[Any]: + if x is None: + out = [] + elif isinstance(x, list): + out = x + elif isinstance(x, (tuple, set)): + out = list(x) + elif hasattr(x, "tolist"): + out = x.tolist() + else: + try: + out = list(x) + except TypeError: + out = [x] + if as_int: + res = [] + for v in out: + try: + res.append(int(v)) + except Exception: + pass + return res + if as_str: + return [str(v) for v in out] + return out + + @staticmethod + def _find_semantic_dir(root: Path, tag_name: str = "SemanticHNMFk") -> Optional[Path]: + root = root.resolve() + candidates = [p for p in root.glob(f"*_{tag_name}") if p.is_dir()] + if candidates: + candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True) + return candidates[0].resolve() + plain = (root / tag_name) + return plain.resolve() if plain.is_dir() else None + + @staticmethod + def _load_pickle(path: Path): + with path.open("rb") as f: + return pickle.load(f) + + @staticmethod + def _load_checkpoint(exp_dir: Path) -> Dict[str, Any]: + """Try common checkpoint names; if not found, search recursively for 'checkpoint*'.""" + exp_dir = exp_dir.resolve() + # Primary expected name + ckpt = (exp_dir / "checkpoint.p").resolve() + candidates: List[Path] = [] + if ckpt.is_file(): + candidates = [ckpt] + else: + alternates = [ + exp_dir / "checkpoint.pkl", + exp_dir / "checkpoint.pickle", + exp_dir / "checkpoint", + ] + candidates = [c.resolve() for c in alternates if c.is_file()] + if not candidates: + globbed = [g.resolve() for g in exp_dir.glob("checkpoint*") if g.is_file()] + if globbed: + candidates = globbed + if not candidates: + # NEW: recursive search fallback + deep = [g.resolve() for g in exp_dir.rglob("checkpoint*") if g.is_file()] + if deep: + # newest first + deep.sort(key=lambda p: p.stat().st_mtime, reverse=True) + candidates = [deep[0]] + + if not candidates: + return {} + + ckpt = candidates[0] + try: + # FIX: correct class reference + return CollectHNMFkLeafBlock._load_pickle(ckpt) + except Exception: + return {} + + @staticmethod + def _safe_rebase(path: str, old_base: Optional[str], new_base: Optional[str]) -> str: + if not path: + return path + p_norm = os.path.normpath(path) + if not old_base or not new_base: + return str(Path(p_norm).resolve()) + try: + old_base_n = os.path.normpath(old_base) + new_base_n = os.path.normpath(new_base) + if p_norm.startswith(old_base_n): + rel = os.path.relpath(p_norm, old_base_n) + return str(Path(os.path.join(new_base_n, rel)).resolve()) + except Exception: + pass + try: + old_seg = re.search(r"(\d+_SemanticHNMFk|SemanticHNMFk)", str(old_base)).group(1) + new_seg = re.search(r"(\d+_SemanticHNMFk|SemanticHNMFk)", str(new_base)).group(1) + return str(Path(p_norm.replace(old_seg, new_seg)).resolve()) + except Exception: + return str(Path(p_norm).resolve()) + + def _log(self, fp: Path, *msgs: str) -> None: + try: + ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + fp.parent.mkdir(parents=True, exist_ok=True) + with fp.open("a", encoding="utf-8") as f: + for m in msgs: + f.write(f"[{ts}] {m}\n") + except Exception: + pass + + # ─────────────────────────── run ──────────────────────────── + + def run(self, bundle: DataBundle) -> None: + df: pd.DataFrame = bundle["df"] + root_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]).expanduser().resolve() + + out_dir = (root_dir / self.tag).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + log_fp = out_dir / "debug.log" + + # Output paths (exact filenames requested) + leaf_data_csv = out_dir / "LEAF_DATA.csv" + leaf_labels_csv = out_dir / "LEAF_LABELS.csv" + summary_txt = out_dir / "summary.txt" # NEW + + # Resolve HNMFk experiment dir — ORDER: bundle → call_settings → auto + tag_name = self.call_settings.get("hnmfk_tag", "SemanticHNMFk") + bundle_key = self.call_settings.get("hnmfk_bundle_key", f"{tag_name}.model_path") + + exp_dir: Optional[Path] = None + src = "" + + # 1) from bundle (model_path preferred, else model object) + if bundle_key in bundle: + exp_dir = Path(bundle[bundle_key]).expanduser().resolve() + src = f"bundle[{bundle_key}]" + elif f"{tag_name}.model" in bundle: + try: + model_obj = bundle[f"{tag_name}.model"] + exp_dir = Path(getattr(model_obj, "experiment_save_path")).expanduser().resolve() + src = f"bundle[{tag_name}.model]" + except Exception: + exp_dir = None + + # 2) explicit in call_settings + if exp_dir is None and self.call_settings.get("hnmfk_dir"): + exp_dir = Path(self.call_settings["hnmfk_dir"]).expanduser().resolve() + src = "call_settings.hnmfk_dir" + + # 3) auto-discover next to SAVE_DIR + if exp_dir is None: + exp_dir = self._find_semantic_dir(root_dir, tag_name=tag_name) + src = f"auto({root_dir})" + + if not exp_dir or not exp_dir.exists(): + self._log(log_fp, f"HNMFk dir not found (src={src}). Writing empty outputs.") + pd.DataFrame(columns=list(df.columns) + ["cluster", "Graph_Name"]).to_csv(leaf_data_csv, index=False, encoding="utf-8-sig") + pd.DataFrame(columns=["Graph_Name", "label", "words"]).to_csv(leaf_labels_csv, index=False, encoding="utf-8-sig") + # NEW: still write an empty summary.txt + with summary_txt.open("w", encoding="utf-8") as f: + f.write("Total documents: 0\nLeaf clusters: 0\n") + self.register_checkpoint("leaf_data_csv", leaf_data_csv) + self.register_checkpoint("leaf_labels_csv", leaf_labels_csv) + bundle[f"{self.tag}.leaf_data_csv"] = leaf_data_csv + bundle[f"{self.tag}.leaf_labels_csv"] = leaf_labels_csv + if self.verbose: + print(f"[{self.tag}] Wrote:\n {leaf_data_csv}\n {leaf_labels_csv}\n {summary_txt}") + return + + ckpt = self._load_checkpoint(exp_dir) + if not ckpt: + self._log(log_fp, f"No checkpoint file in {exp_dir}. Writing empty outputs.") + pd.DataFrame(columns=list(df.columns) + ["cluster", "Graph_Name"]).to_csv(leaf_data_csv, index=False, encoding="utf-8-sig") + pd.DataFrame(columns=["Graph_Name", "label", "words"]).to_csv(leaf_labels_csv, index=False, encoding="utf-8-sig") + with summary_txt.open("w", encoding="utf-8") as f: + f.write("Total documents: 0\nLeaf clusters: 0\n") + self.register_checkpoint("leaf_data_csv", leaf_data_csv) + self.register_checkpoint("leaf_labels_csv", leaf_labels_csv) + bundle[f"{self.tag}.leaf_data_csv"] = leaf_data_csv + bundle[f"{self.tag}.leaf_labels_csv"] = leaf_labels_csv + if self.verbose: + print(f"[{self.tag}] Wrote:\n {leaf_data_csv}\n {leaf_labels_csv}\n {summary_txt}") + return + + node_save_paths: Dict[str, str] = ckpt.get("node_save_paths", {}) or {} + root_name: str = ckpt.get("root_name") or "Root" + old_base = ckpt.get("experiment_name") or ckpt.get("experiment_save_path") + new_base = str(exp_dir) + + final_df_parts: List[pd.DataFrame] = [] + label_parts: List[pd.DataFrame] = [] + visited = 0 + leaves_seen = 0 + + def _resolve(p: str) -> str: + return self._safe_rebase(p, old_base, new_base) + + def _load_node_safely(path: Path): + try: + return self._load_pickle(path) + except Exception as e: + self._log(log_fp, f"Failed to load node: {path} :: {e}") + return None + + def visit(node_name: str): + nonlocal visited, leaves_seen + if node_name not in node_save_paths: + self._log(log_fp, f"node_save_paths missing entry for '{node_name}'") + return + node_path = Path(_resolve(node_save_paths[node_name])).expanduser().resolve() + visited += 1 + if not node_path.exists(): + self._log(log_fp, f"node pickle missing: {node_path}") + return + + node_obj = _load_node_safely(node_path) + if node_obj is None: + return + + # children + for cname in self._to_list(getattr(node_obj, "child_node_names", []), as_str=True): + visit(cname) + + # leaf? + if not bool(getattr(node_obj, "leaf", False)): + return + + leaves_seen += 1 + original_indices = self._to_list(getattr(node_obj, "original_indices", None), as_int=True) + if not original_indices: + self._log(log_fp, f"leaf with empty original_indices: {node_path}") + return + + # extract ALL columns for those rows + try: + node_df = df.iloc[original_indices].copy() + except Exception as e: + self._log(log_fp, f"iloc failed for indices (len={len(original_indices)}): {e}") + return + + # num_clusters from W or signature + W = getattr(node_obj, "W", None) + if W is None: + sig = getattr(node_obj, "signature", None) + if sig is None: + num_clusters = 1 + else: + sig = np.asarray(sig) + if sig.ndim == 1: + sig = sig.reshape(-1, 1) + num_clusters = int(sig.shape[1]) + else: + W = np.asarray(W) + if W.ndim == 1: + W = W.reshape(-1, 1) + num_clusters = int(W.shape[1]) + + node_dir = node_path.parent + cluster_csv = node_dir / f"cluster_for_k={num_clusters}.csv" + cluster_membership = None + if cluster_csv.is_file(): + try: + cluster_membership = pd.read_csv(cluster_csv) + except Exception as e: + self._log(log_fp, f"read_csv failed: {cluster_csv} :: {e}") + + # attach cluster column + if isinstance(cluster_membership, pd.DataFrame) and "cluster" in cluster_membership.columns: + if len(cluster_membership) == len(node_df): + node_df["cluster"] = cluster_membership["cluster"].to_numpy() + else: + # try to align if cluster CSV has a 'doc_index' column + if "doc_index" in cluster_membership.columns: + try: + aligned = cluster_membership.set_index("doc_index").loc[original_indices] + node_df["cluster"] = aligned["cluster"].to_numpy() + except Exception: + m = min(len(cluster_membership), len(node_df)) + node_df = node_df.iloc[:m].copy() + node_df["cluster"] = cluster_membership["cluster"].iloc[:m].to_numpy() + self._log(log_fp, f"cluster length mismatch; truncated to {m} :: {cluster_csv}") + else: + m = min(len(cluster_membership), len(node_df)) + node_df = node_df.iloc[:m].copy() + node_df["cluster"] = cluster_membership["cluster"].iloc[:m].to_numpy() + self._log(log_fp, f"cluster length mismatch; truncated to {m} :: {cluster_csv}") + else: + node_df["cluster"] = 0 # default single cluster + + # Graph_Name from parent dir of node pickle + cluster id + graph_name_part = os.path.basename(os.path.dirname(str(node_path))) + for k_val in sorted(pd.unique(node_df["cluster"])): + sub = node_df[node_df["cluster"] == k_val].copy() + sub["Graph_Name"] = f"{graph_name_part}_{k_val}" + final_df_parts.append(sub) + + # labels + cs_fp = node_dir / "cluster_summaries.csv" + tw_fp = node_dir / "top_words.csv" + if cs_fp.is_file() and tw_fp.is_file(): + try: + cs = pd.read_csv(cs_fp) + tw = pd.read_csv(tw_fp) + words_map: Dict[str, str] = {} + for col in tw.columns: + words_map[str(col)] = ",".join(tw[col].dropna().astype(str).tolist()) + cs = cs.copy() + cs["cluster"] = cs["cluster"].astype(str) + cs["words"] = cs["cluster"].map(words_map).fillna("") + cs["Graph_Name"] = cs["cluster"].apply(lambda c: f"{graph_name_part}_{c}") + label_parts.append(cs[["Graph_Name", "label", "words"]]) + except Exception as e: + self._log(log_fp, f"label build failed at {node_dir}: {e}") + + # Walk from root + if root_name not in node_save_paths: + self._log(log_fp, f"root_name '{root_name}' missing in node_save_paths; nothing to traverse.") + else: + visit(root_name) + + # Concatenate & write (always write something) + if final_df_parts: + final_df = pd.concat(final_df_parts, ignore_index=True) + else: + final_df = pd.DataFrame(columns=list(df.columns) + ["cluster", "Graph_Name"]) + final_df.to_csv(leaf_data_csv, index=False, encoding="utf-8-sig") + + if label_parts: + labels_df = pd.concat(label_parts, ignore_index=True) + else: + labels_df = pd.DataFrame(columns=["Graph_Name", "label", "words"]) + labels_df.to_csv(leaf_labels_csv, index=False, encoding="utf-8-sig") + + # NEW: Write summary.txt (totals + per-cluster counts) + try: + if not final_df.empty and "Graph_Name" in final_df.columns: + counts = final_df.groupby("Graph_Name").size().sort_values(ascending=False) + else: + counts = pd.Series(dtype=int) + total_docs = int(len(final_df)) + num_leaf_clusters = int(len(counts)) + lines = [ + f"Total documents: {total_docs}", + f"Leaf clusters: {num_leaf_clusters}", + "", + ] + lines += [f"{name}\t{int(cnt)}" for name, cnt in counts.items()] + with summary_txt.open("w", encoding="utf-8") as f: + f.write("\n".join(lines)) + except Exception as e: + self._log(log_fp, f"Failed to write summary.txt: {e}") + + # Checkpoint + bundle exposure + self.register_checkpoint("df", leaf_data_csv) + self.register_checkpoint("leaf_labels_csv", leaf_labels_csv) + bundle[f"{self.tag}.df"] = leaf_data_csv + bundle[f"{self.tag}.leaf_labels_csv"] = leaf_labels_csv + # Optional: also expose summary path in the bundle (no need to register) + bundle[f"{self.tag}.summary_txt"] = summary_txt # NEW + + # Summary in logs + stdout (if verbose) + self._log( + log_fp, + f"exp_dir={exp_dir}", + f"nodes_seen={visited}, leaves_seen={leaves_seen}", + f"wrote LEAF_DATA rows={len(final_df)}", + f"wrote LEAF_LABELS rows={len(labels_df)}", + f"wrote summary at {summary_txt}", + ) + if self.verbose: + print(f"[{self.tag}] Wrote:\n {leaf_data_csv}\n {leaf_labels_csv}\n {summary_txt}") \ No newline at end of file diff --git a/TELF/pipeline/blocks/ocelot_filter_block.py b/TELF/pipeline/blocks/ocelot_filter_block.py old mode 100755 new mode 100644 diff --git a/TELF/pipeline/blocks/orca_block.py b/TELF/pipeline/blocks/orca_block.py index 343ee3d7..0039a8a1 100644 --- a/TELF/pipeline/blocks/orca_block.py +++ b/TELF/pipeline/blocks/orca_block.py @@ -1,5 +1,6 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Tuple, Optional, Sequence +import pandas as pd from .base_block import AnimalBlock from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY @@ -49,12 +50,24 @@ def run(self, bundle: DataBundle) -> None: df = drop_columns_if_exist(df, cols = ['slic_affiliations', 'slic_author_ids', 'slic_authors']) orca = Orca(**self.init_settings) - df = clean_affiliations(df) + + for col in ["affiliations", "year"]: + if col not in df.columns: + df[col] = pd.NA + + df = clean_affiliations(df) orca_map_df = orca.run(df) df = clean_affiliations(df) + if df.get('affiliations', pd.Series(dtype=object)).notna().any(): + df = clean_affiliations(df) + orca_map_df = orca.run(df) - orca_map_df = orca_map_df.dropna(subset=['scopus_ids']).reset_index(drop=True) + # orca_map_df = orca_map_df.dropna(subset=['scopus_ids']).reset_index(drop=True) + if orca_map_df.get('scopus_ids', pd.Series(dtype=object)).notna().any(): + orca_map_df = orca_map_df.dropna(subset=['scopus_ids']).reset_index(drop=True) + elif orca_map_df.get('s2_ids', pd.Series(dtype=object)).notna().any(): + orca_map_df = orca_map_df.dropna(subset=['s2_ids']).reset_index(drop=True) orca_map_df = add_num_known_col(orca_map_df) orca_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag @@ -62,9 +75,12 @@ def run(self, bundle: DataBundle) -> None: df = orca.apply(df) - df = clean_affiliations(df) - - df = prep_affiliations(df) + # df = clean_affiliations(df) + # df = prep_affiliations(df) + if df.get('slic_affiliations', pd.Series(dtype=object)).notna().any(): + df = clean_affiliations(df) + df = prep_affiliations(df) + if 'type' in df.columns: orca_summary_path = save_path=orca_dir / 'type_summary.csv' diff --git a/TELF/pipeline/blocks/peacock_stats_block.py b/TELF/pipeline/blocks/peacock_stats_block.py index 8a290e0f..12635563 100644 --- a/TELF/pipeline/blocks/peacock_stats_block.py +++ b/TELF/pipeline/blocks/peacock_stats_block.py @@ -1,239 +1,126 @@ # pipeline/blocks/peacock_stats_block.py from __future__ import annotations from pathlib import Path -from typing import Any, Dict, Sequence, Optional, Tuple - +from typing import Any, Dict, Optional, Sequence, Tuple, Literal import pandas as pd from .base_block import AnimalBlock -from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY, RESULTS_DEFAULT -# Peacock aggregation function -from ...post_processing.Peacock.Utility import aggregate_ostats +from .block_helpers.peacock_renderer import PeacockRenderer +from .block_helpers.hnmfk_paths import NodeSource, HNMFkNodeSource, GlobNodeSource -# Peacock plotting functions (the ones you pasted from Plot/plot.py) -from ...post_processing.Peacock.Plot import ( - plot_heatmap, - plot_bar, - plot_scatter, -) -# Convenience aliases -plot_hist = plot_bar -plot_scatter3D = plot_scatter +Mode = Literal["single", "hnmfk", "glob"] class PeacockStatsBlock(AnimalBlock): - CANONICAL_NEEDS: Tuple[str, ...] = ("df", SAVE_DIR_BUNDLE_KEY) + CANONICAL_NEEDS: Tuple[str, ...] = ("df", ) def __init__( self, *, needs: Sequence[str] = CANONICAL_NEEDS, provides: Sequence[str] = ("outpath",), + mode: Mode = "single", + # renderer args hist_stats: Sequence[str] = ("paper_count", "num_citations"), hist_ylabels: Optional[Dict[str, str]] = None, col_names: Optional[Dict[str, str]] = None, affiliation_palette: Optional[Dict[str, str]] = None, country: Optional[str] = None, + cluster_col: Optional[str] = "cluster", # NEW + # discovery args + experiment_path: Optional[str] = None, # for mode="hnmfk" + glob_root: Optional[str] = None, # for mode="glob" + skip_completed: bool = True, **kw: Any, ) -> None: + self.mode = mode + self.experiment_path = Path(experiment_path).expanduser().resolve() if experiment_path else None + self.glob_root = Path(glob_root).expanduser().resolve() if glob_root else None + self.skip_completed = skip_completed + + self.renderer = PeacockRenderer( + hist_stats=hist_stats, + hist_ylabels=hist_ylabels, + col_names=col_names, + affiliation_palette=affiliation_palette, + country=country, + cluster_col=cluster_col, + ) + + conds = list(kw.pop("conditional_needs", ())) + if self.mode == "hnmfk" and not self.experiment_path: + conds.append(("model_path", lambda _b, _s: True)) + super().__init__( needs=needs, provides=provides, tag="PeacockStats", init_settings={}, call_settings={}, + conditional_needs=tuple(conds), **kw, ) - self.hist_stats = tuple(hist_stats) - self.hist_ylabels = hist_ylabels or { - "paper_count": "Number of Papers", - "num_citations": "Number of Citations", - "attribution_percentage": "Attribution Percentage", - } - self.col_names = col_names or { - "id": "eid", - "authors": "slic_authors", - "author_ids": "slic_author_ids", - "affiliations": "slic_affiliations", - "funding": "funding", - "citations": "num_citations", - "references": "references", - } - self.affiliation_palette = affiliation_palette or {} - self.country = country - def run(self, bundle: DataBundle) -> None: - # 1) Inputs & cleanup - df: pd.DataFrame = bundle["df"] - out_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]) - out_dir.mkdir(parents=True, exist_ok=True) - - # ensure affiliation column is string - aff_col = self.col_names["affiliations"] - df[aff_col] = df[aff_col].apply(lambda x: x if isinstance(x, str) else str(x)) - - df = ( - df - .dropna(subset=[ - self.col_names["id"], - self.col_names["authors"], - self.col_names["author_ids"], - aff_col, - ]) - .assign(year=lambda d: d.year.astype(int)) - ) + def _source(self, bundle: DataBundle) -> Optional[NodeSource]: + if self.mode == "hnmfk": + exp = self.experiment_path or Path(str(bundle["model_path"])) + return HNMFkNodeSource(exp) + if self.mode == "glob": + root = self.glob_root or Path(bundle.get(SAVE_DIR_BUNDLE_KEY, RESULTS_DEFAULT)) + return GlobNodeSource(root) + return None - filters = {"country": self.country} if self.country else None - - # 2) Write top‐100 CSVs - author_stats = aggregate_ostats(df, key="author_id", top_n=100, col_names=self.col_names, filters=filters, by_year=False) - affiliation_stats = aggregate_ostats(df, key="affiliation_id", top_n=100, col_names=self.col_names, filters=filters, by_year=False) - author_stats.to_csv(out_dir / "top_authors.csv", index=False) - affiliation_stats.to_csv(out_dir / "top_affiliations.csv", index=False) - - # 3) Shared “top‐10 by citations” args - auth_args = dict( - key="author_id", - top_n=10, - sort_by="num_citations", - col_names=self.col_names, - by_year=True, - filters=filters, - ) - aff_args = dict( - key="affiliation_id", - top_n=10, - sort_by="num_citations", - col_names=self.col_names, - by_year=True, - filters=filters, - ) - - # 4) Heatmaps — pivot then call plot_heatmap - # ─ Authors, citations - auth_heat = aggregate_ostats(df, **auth_args) - pivot_c = auth_heat.pivot(index="year", columns="author", values="num_citations").fillna(0) - plot_heatmap( - pivot_c, - cmap="jet", - interpolation="gaussian", - fname=str(out_dir / "author_heatmap_citations.png"), - interactive=False, - title="Author Citations by Year", - xlabel="Author", - ylabel="Year", - ) - # ─ Authors, papers - pivot_p = auth_heat.pivot(index="year", columns="author", values="paper_count").fillna(0) - plot_heatmap( - pivot_p, - cmap="jet", - interpolation="gaussian", - fname=str(out_dir / "author_heatmap_papers.png"), - interactive=False, - title="Author Papers by Year", - xlabel="Author", - ylabel="Year", - ) + def _parent_out_dir(self, bundle: DataBundle) -> Path: + base = Path(bundle.get(SAVE_DIR_BUNDLE_KEY, RESULTS_DEFAULT)) + return base / self.tag - # ─ Affiliations, citations - aff_heat = aggregate_ostats(df, **aff_args) - pivot_c2 = aff_heat.pivot(index="year", columns="affiliation", values="num_citations").fillna(0) - plot_heatmap( - pivot_c2, - cmap="jet", - interpolation="gaussian", - fname=str(out_dir / "affiliation_heatmap_citations.png"), - interactive=False, - title="Affiliation Citations by Year", - xlabel="Affiliation", - ylabel="Year", - ) - # ─ Affiliations, papers - pivot_p2 = aff_heat.pivot(index="year", columns="affiliation", values="paper_count").fillna(0) - plot_heatmap( - pivot_p2, - cmap="jet", - interpolation="gaussian", - fname=str(out_dir / "affiliation_heatmap_papers.png"), - interactive=False, - title="Affiliation Papers by Year", - xlabel="Affiliation", - ylabel="Year", - ) - - # 5) Histograms (bar‐plots) - auth_hist = aggregate_ostats(df, **{**auth_args, "by_year": False}) - plot_hist( - auth_hist, - x="author", - ys=list(self.hist_stats), - fname=str(out_dir / "author_hist.png"), - interactive=False, - # cmap="husl", # ← remove this line... - title="Author Statistics Histogram", - xlabel="Author", - ylabel=self.hist_ylabels[self.hist_stats[0]], - ) - aff_hist = aggregate_ostats(df, **{**aff_args, "by_year": False}) - plot_hist( - aff_hist, - x="affiliation", - ys=list(self.hist_stats), - fname=str(out_dir / "affiliation_hist.png"), - interactive=False, - # cmap="husl", - title="Affiliation Statistics Histogram", - xlabel="Affiliation", - ylabel=self.hist_ylabels[self.hist_stats[0]], - ) - - # 6) 3-D scatter - plot_scatter3D( - df, - x="paper_count", - y="attribution_percentage", - z="num_citations", - agg_func=aggregate_ostats, - agg_kwargs=auth_args, - fname=str(out_dir / "author_scatter.png"), - interactive=False, - log_z=True, - hue="affiliation", - labels="author", - title="Author Stats Scatter3D", - xlabel="Paper Count", - ylabel="Attribution Percentage", - zlabel="Num. Citations", - base_palette=self.affiliation_palette, - ) - plot_scatter3D( - df, - x="paper_count", - y="attribution_percentage", - z="num_citations", - agg_func=aggregate_ostats, - agg_kwargs=aff_args, - fname=str(out_dir / "affiliation_scatter.png"), - interactive=False, - log_z=True, - hue="country", - labels="affiliation", - title="Affiliation Stats Scatter3D", - xlabel="Paper Count", - ylabel="Attribution Percentage", - zlabel="Num. Citations", - base_palette=self.affiliation_palette, - ) - - # 7) Dummy checkpoint under your single `provides` key - if SAVE_DIR_BUNDLE_KEY in bundle: - ckpt_dir = out_dir / self.tag - ckpt_dir.mkdir(parents=True, exist_ok=True) - final_csv = ckpt_dir / "none.csv" - final_csv.write_text("") # empty placeholder - self.register_checkpoint(self.provides[0], final_csv) - bundle[f"{self.tag}.{self.provides[0]}"] = final_csv - # (no return: AnimalBlock.__call__ returns the bundle) + def run(self, bundle: DataBundle) -> None: + # single mode: unchanged + if self.mode == "single": + df: pd.DataFrame = bundle["df"] + out_dir = Path(bundle.get(SAVE_DIR_BUNDLE_KEY, RESULTS_DEFAULT)) + self.renderer.render(df, out_dir) + ckpt_dir = self._parent_out_dir(bundle); ckpt_dir.mkdir(parents=True, exist_ok=True) + marker = ckpt_dir / "none.csv"; marker.write_text("status\nok\n") + self.register_checkpoint(self.provides[0], marker) + bundle[f"{self.tag}.{self.provides[0]}"] = marker + return + + # per-node modes (hnmfk or glob) + source = self._source(bundle) + assert source is not None, "Invalid configuration: per-node mode requires a NodeSource." + + produced = [] + for node in source.iter_nodes(): + if self.skip_completed and (node.dir / "peacock" / "PeacockStats.done").exists(): + produced.append(node.dir / "peacock") + continue + + if not node.csv.exists(): + continue + + df_local = pd.read_csv(node.csv) + + # If the dataframe has clusters, split them into subfolders + if "cluster" in df_local.columns: + for cid, df_c in df_local.groupby("cluster", dropna=False): + safe = "nan" if pd.isna(cid) else str(cid).replace("/", "_") + out_dir = node.dir / str(safe) / "peacock" + self.renderer.render(df_c, out_dir) + (out_dir / "PeacockStats.done").write_text("ok") + produced.append(out_dir) + else: + out_dir = node.dir / "peacock" + self.renderer.render(df_local, out_dir) + (out_dir / "PeacockStats.done").write_text("ok") + produced.append(out_dir) + + # block-level checkpoint + ckpt_dir = self._parent_out_dir(bundle); ckpt_dir.mkdir(parents=True, exist_ok=True) + registry = ckpt_dir / "per_node_registry.txt" + registry.write_text("\n".join(map(str, produced))) + self.register_checkpoint(self.provides[0], registry) + bundle[f"{self.tag}.{self.provides[0]}"] = registry diff --git a/TELF/pipeline/blocks/sbatch_block.py b/TELF/pipeline/blocks/sbatch_block.py index a349dcde..8a58e407 100644 --- a/TELF/pipeline/blocks/sbatch_block.py +++ b/TELF/pipeline/blocks/sbatch_block.py @@ -35,13 +35,13 @@ class SBatchBlock(AnimalBlock): Behaviour --------- - • **Skip submission** if every declared checkpoint already exists. + • **Skip submission** if every declared checkpoint already exists. • Otherwise: 1. create a staging dir `/` (spaces → `_`); 2. serialize the current bundle + wrapped block there; - 3. write a `run_block.py` runner (faulthandler enabled); - 4. write a hardened `run.sbatch` (threads pinned, Arrow off); - 5. `sbatch run.sbatch` and `sys.exit` so the pipeline resumes later. + 3. write a `run_block.py` runner (faulthandler + hardened jsonpickle); + 4. write a hardened `run.sbatch` (threads pinned, Arrow off, PYTHONPATH set); + 5. `sbatch run.sbatch` and exit so the pipeline resumes later. """ def __init__( self, @@ -58,7 +58,7 @@ def __init__( call_settings: Dict[str, Any] | None = None, **kw, ) -> None: - tag=wrapped_block.tag + tag = wrapped_block.tag super().__init__( needs=needs or wrapped_block.needs, provides=provides or wrapped_block.provides, @@ -101,7 +101,7 @@ def run(self, bundle: DataBundle) -> None: done = workdir / "_complete.json" if done.exists(): # runner wrote out a small dict: { provide_name: value, ... } - resumed: dict = jsonpickle.decode(done.read_text()) + resumed: dict = jsonpickle.decode(done.read_text(), keys=True) wrapped_tag = self.wrapped_block.tag for p in self.provides: ns = f"{wrapped_tag}.{p}" @@ -112,33 +112,113 @@ def run(self, bundle: DataBundle) -> None: # ensure SAVE_DIR is absolute bundle[SAVE_DIR_BUNDLE_KEY] = str(base_dir) - # 3) serialize bundle + block + # 3) serialize bundle + block (keys=True preserves non-string dict keys) (workdir / "input_bundle.json").write_text( - jsonpickle.encode(bundle), "utf-8" + jsonpickle.encode(bundle, keys=True), encoding="utf-8" ) (workdir / "block.json").write_text( - jsonpickle.encode(self.wrapped_block), "utf-8" + jsonpickle.encode(self.wrapped_block, keys=True), encoding="utf-8" ) + # sidecar with import path for robust re-instantiation + block_cls_path = f"{self.wrapped_block.__class__.__module__}:{self.wrapped_block.__class__.__qualname__}" + (workdir / "block_class.txt").write_text(block_cls_path, encoding="utf-8") + + # 4) runner script (Template placeholders; adds sys.path; spaCy Table handler; dict→class fallback) + from string import Template - # 4) runner script runner = workdir / "run_block.py" - runner.write_text(textwrap.dedent(f""" - import faulthandler, jsonpickle - faulthandler.enable() + project_root = Path.cwd() # adjust if your repo root differs at submit time + + tpl = Template(""" +import sys, importlib, types, faulthandler, jsonpickle +faulthandler.enable() + +# Make sure we can import your package/modules +sys.path.insert(0, $PROJECT_ROOT) + +# ---- jsonpickle handler: spaCy Table needs constructor run before inserts ---- +from jsonpickle.handlers import BaseHandler +try: + from spacy.lookups import Table as _SpaCyTable + _TABLES = [_SpaCyTable] +except Exception: + _TABLES = [] + +class _TableHandler(BaseHandler): + def flatten(self, obj, data): + data['py/object'] = obj.__class__.__module__ + '.' + obj.__class__.__name__ + data['items'] = list(obj.items()) + return data + + def restore(self, data): + cls = self.restore_class(data) + inst = cls() # ensure __init__ runs (creates .bloom, etc.) + if 'items' in data: + kv = data['items'] + else: + state = data.get('py/state', {}) + if isinstance(state, dict): + mapping = state.get('data', state) + else: + mapping = {} + kv = list(getattr(mapping, 'items', lambda: [])()) + if kv: + inst.update(dict(kv)) + return inst - bundle = jsonpickle.decode(open(r"{workdir/'input_bundle.json'}").read()) - block = jsonpickle.decode(open(r"{workdir/'block.json'}").read()) - out_bundle = block(bundle) +for _T in _TABLES: + jsonpickle.handlers.register(_T, _TableHandler) - # collect just the wrapped block's provides - result = {{}} - for p in {list(self.provides)!r}: - result[p] = out_bundle[f"{{block.tag}}.{{p}}"] +# ---- decode inputs (keys=True for non-string keys elsewhere) ---- +with open($INPUT_JSON, "r", encoding="utf-8") as _f: + bundle = jsonpickle.decode(_f.read(), keys=True) +with open($BLOCK_JSON, "r", encoding="utf-8") as _f: + block = jsonpickle.decode(_f.read(), keys=True) - open(r"{done}", "w").write(jsonpickle.encode(result)) - """).strip(), "utf-8") +# If import failed and we got a dict, rebuild the class instance using the sidecar +if isinstance(block, dict): + with open($BLOCK_CLASS_TXT, "r", encoding="utf-8") as _f: + cls_path = _f.read().strip() + mod_name, qualname = cls_path.split(":", 1) + mod = importlib.import_module(mod_name) + cls = mod + for part in qualname.split("."): + cls = getattr(cls, part) + # instantiate and copy state (skip framework-managed attrs) + NON_STATE = {"tag", "needs", "provides", "conditional_needs", "load_checkpoint"} + inst = cls() + for k, v in block.items(): + if k not in NON_STATE: + setattr(inst, k, v) + # keep tag/needs/provides if present in the dict + for k in ("tag", "needs", "provides", "conditional_needs"): + if k in block: + setattr(inst, k, block[k]) + block = inst + +# Run the block +out_bundle = block(bundle) + +# collect just the wrapped block's provides +result = {} +for p in $PROVIDES: + result[p] = out_bundle[block.tag + "." + p] + +with open($DONE_JSON, "w", encoding="utf-8") as _f: + _f.write(jsonpickle.encode(result, keys=True)) +""") + + script = tpl.substitute( + PROJECT_ROOT=repr(str(project_root)), + INPUT_JSON=repr(str(workdir / "input_bundle.json")), + BLOCK_JSON=repr(str(workdir / "block.json")), + BLOCK_CLASS_TXT=repr(str(workdir / "block_class.txt")), + DONE_JSON=repr(str(done)), + PROVIDES=repr(list(self.provides)), + ) + runner.write_text(textwrap.dedent(script), encoding="utf-8") - # 5) sbatch script + # 5) sbatch script (export PYTHONPATH to make imports resolvable in the job) sbatch = workdir / "run.sbatch" lines = [ "#!/bin/bash", @@ -155,11 +235,12 @@ def run(self, bundle: DataBundle) -> None: "export PANDAS_ARROW_DISABLED=1", "", f"cd {shquote(workdir)}", + f'export PYTHONPATH={shquote(project_root)}:"$PYTHONPATH"', ] match self.venv_type: case "venv": - lines.append(f"source {shquote(Path(self.venv_path).resolve()/ 'bin'/'activate')}") + lines.append(f"source {shquote(Path(self.venv_path).resolve() / 'bin' / 'activate')}") case "conda": lines += [ "source $(conda info --base)/etc/profile.d/conda.sh", @@ -174,7 +255,7 @@ def run(self, bundle: DataBundle) -> None: raise ValueError(f"Unknown venv_type {self.venv_type!r}") lines.append(f"python -Xfaulthandler {shquote(runner)}") - sbatch.write_text("\n".join(lines), "utf-8") + sbatch.write_text("\n".join(lines), encoding="utf-8") # 6) submit and exit print(f"🚀 {self.tag}: submitting via sbatch …") diff --git a/TELF/pipeline/blocks/semantic_hnmfk_block.py b/TELF/pipeline/blocks/semantic_hnmfk_block.py index eea6412f..d842f170 100644 --- a/TELF/pipeline/blocks/semantic_hnmfk_block.py +++ b/TELF/pipeline/blocks/semantic_hnmfk_block.py @@ -7,6 +7,8 @@ import numpy as np import pandas as pd import scipy.sparse as ss +import traceback +import warnings from ...factorization import HNMFk from ...pre_processing import Beaver @@ -55,12 +57,24 @@ def __call__(self, original_idx: np.ndarray): save_path=None, ) return X.T.tocsr(), {"vocab": vocab} - except Exception: - return ss.csr_matrix([[1]]), { - "stop_reason": "documents_words could not build matrix", + # except Exception: + # return ss.csr_matrix([[1]]), { + # "stop_reason": "documents_words could not build matrix", + # } + except Exception as e: + tb = traceback.format_exc() + warnings.warn(f"documents_words failed on {len(original_idx)} docs: {e}\n{tb}") + + # Preserve the docs dimension so you can see it's not 1 + n_docs = int(len(original_idx)) + placeholder = ss.csr_matrix((1, n_docs)) # 1 feature × N docs for H-mode + return placeholder, { + "stop_reason": "documents_words failed", + "exception": repr(e), + "traceback": tb, + "n_docs": n_docs, } - # ------------------------------------------------------------------ # # Main block # # ------------------------------------------------------------------ # diff --git a/TELF/pipeline/blocks/spacey_NER_block.py b/TELF/pipeline/blocks/spacey_NER_block.py new file mode 100644 index 00000000..8c70e955 --- /dev/null +++ b/TELF/pipeline/blocks/spacey_NER_block.py @@ -0,0 +1,179 @@ +from pathlib import Path +from typing import Any, Dict, List, Optional +import json +import pandas as pd + +from .base_block import AnimalBlock +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY + +class SpacyNERBlock(AnimalBlock): + """ + spaCy NER over one or more text columns (default: ['title', 'abstract']). + + Adds ONE new column to the DataFrame: + - `output_column` (default: 'ner_by_label'): a STRING of a dictionary that + can be round-tripped with `ast.literal_eval`. The dict maps: + { : [unique entity strings for the row] } + + Artifacts: + - //.csv (enriched DataFrame) + + Requirements: + - spaCy with a model installed (default: 'en_core_web_lg'). + + Parameters + ---------- + needs : tuple + Bundle keys to read. Default: ("df",). + provides : tuple + Bundle keys to write. Default: ("df",). + text_columns : Optional[List[str]] + Which DF columns to run NER on. Default: ["title", "abstract"]. + id_field : str + (Unused for output but kept for compatibility.) Default: "eid". + spacy_model : str + spaCy model name to load. Default: "en_core_web_lg". + batch_size : int + spaCy pipe batch size. Default: 256. + n_process : int + Number of processes for spaCy pipe. Default: 1. + drop_existing : bool + If True, drop an existing output column before writing. Default: True. + output_column : str + Name of the new column to write. Default: "ner_by_label". + tag : str + Block tag and artifact folder name. Default: "spaceyNER". + init_settings : Optional[Dict[str, Any]] + Extra init settings; merged into defaults. + """ + + CANONICAL_NEEDS = ("df",) + + def __init__( + self, + *, + needs=CANONICAL_NEEDS, + provides=("df",), + text_columns: Optional[List[str]] = None, + id_field: str = "eid", + spacy_model: str = "en_core_web_lg", + batch_size: int = 256, + n_process: int = 1, + drop_existing: bool = True, + output_column: str = "ner_by_label", + tag: str = "spaceyNER", + init_settings: Optional[Dict[str, Any]] = None, + **kw, + ): + self.id_field = id_field + + default_init = { + "text_columns": text_columns or ["title", "abstract"], + "spacy_model": spacy_model, + "batch_size": int(batch_size), + "n_process": int(n_process), + "drop_existing": bool(drop_existing), + "output_column": output_column, + "verbose": True, + } + + super().__init__( + needs=needs, + provides=provides, + tag=tag, + init_settings=self._merge(default_init, init_settings), + call_settings={}, # not used here + **kw, + ) + + # ------------------------------- helpers + + def _load_spacy(self): + try: + import spacy # local import so module is optional until needed + except Exception as e: + raise RuntimeError( + f"[{self.tag}] spaCy is not installed. Please install spaCy and a model." + ) from e + + model = self.init_settings["spacy_model"] + try: + # Disable components we don't need for speed. + nlp = spacy.load(model, disable=["tagger", "lemmatizer", "textcat"]) + except OSError as e: + raise RuntimeError( + f"[{self.tag}] spaCy model '{model}' is not installed.\n" + f"Install it with: python -m spacy download {model}" + ) from e + return nlp + + # ------------------------------- run + + def run(self, bundle: DataBundle) -> None: + # 1) Load input DF + df: pd.DataFrame = self.load_path(bundle[self.needs[0]]) + print(f"\n[{self.tag}] ====================================================") + print(f"[{self.tag}] df.shape = {df.shape}") + + text_cols: List[str] = list(self.init_settings.get("text_columns", ["title", "abstract"])) + present = [c for c in text_cols if c in df.columns] + missing = [c for c in text_cols if c not in df.columns] + if not present: + raise ValueError(f"[{self.tag}] None of the requested text columns are present: {text_cols}") + if missing: + print(f"[{self.tag}] WARNING: missing text columns {missing}; will skip them") + + # 2) Output dir + root = Path(bundle.get(SAVE_DIR_BUNDLE_KEY, ".")) + out = root / self.tag + out.mkdir(parents=True, exist_ok=True) + print(f"[{self.tag}] output dir = {out}") + + # 3) Load spaCy + nlp = self._load_spacy() + batch_size = int(self.init_settings.get("batch_size", 256)) + n_process = int(self.init_settings.get("n_process", 1)) + drop_existing = bool(self.init_settings.get("drop_existing", True)) + out_col = str(self.init_settings.get("output_column", "ner_by_label")) + + # 4) Prepare DF; optionally drop existing output column + df_proc = df.copy() + if drop_existing and out_col in df_proc.columns: + df_proc = df_proc.drop(columns=[out_col]) + print(f"[{self.tag}] dropped existing column: {out_col}") + + # We'll aggregate per-row across ALL present text columns. + # For each row, maintain a dict[label] -> list[str] (unique, preserve insertion order) + num_rows = len(df_proc) + agg_by_row: List[Dict[str, List[str]]] = [dict() for _ in range(num_rows)] + + # 5) Run NER column-by-column and merge results per row + for col in present: + print(f"[{self.tag}] NER on column = {col}") + texts = df_proc[col].fillna("").astype(str).tolist() + + for i, doc in enumerate(nlp.pipe(texts, batch_size=batch_size, n_process=n_process)): + if not doc.ents: + continue + row_dict = agg_by_row[i] + for ent in doc.ents: + label = ent.label_ + text = ent.text + lst = row_dict.setdefault(label, []) + # preserve insertion order & uniqueness + if text not in lst: + lst.append(text) + + # 6) Serialize each row's dict to a string compatible with ast.literal_eval + # Using JSON ensures safe, unambiguous formatting; ast.literal_eval accepts JSON literals. + serialized = [json.dumps(d, ensure_ascii=False) for d in agg_by_row] + + df_proc[out_col] = serialized + + # 7) Save artifacts + register + df_path = out / f"{self.tag}.csv" + df_proc.to_csv(df_path, index=False, encoding="utf-8-sig") + self.register_checkpoint(self.provides[0], df_path) + bundle[f"{self.tag}.{self.provides[0]}"] = df_proc + print(f"[{self.tag}] saved df → {df_path}") + print(f"[{self.tag}] added column = {out_col}") \ No newline at end of file diff --git a/TELF/pipeline/blocks/termite_neo4j_block.py b/TELF/pipeline/blocks/termite_neo4j_block.py new file mode 100644 index 00000000..c63ed74c --- /dev/null +++ b/TELF/pipeline/blocks/termite_neo4j_block.py @@ -0,0 +1,1077 @@ +# TELF/pipeline/blocks/termite_neo4j_block.py +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Dict, Sequence, Tuple, Optional, List, Callable +import pandas as pd +from copy import deepcopy +import ast +import json +import re +import yaml + +from .base_block import AnimalBlock +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY + +# --- Termite + constants --- +from TELF.applications import Termite +from TELF.applications.Termite.neo4j_termite import ( + ENTITY, RETURN_TYPE, ATTRIBUTES, ET, YEAR_TYPE, FROM_COL, DOCUMENT_TYPE, + ROW_INDEX, ATTR_COL, ATTR_NAME, TT, R, DOCUMENT_YEAR_RELATION, HT, + AUTHOR_DOCUMENT_RELATION, AUTHOR_ID_TYPE, EXTRACT_H, DOCUMENT_CITES_RELATION, + DOCUMENT_CITED_RELATION, DOCUMENT_TYPE_SCOPUS, EXTRACT_T, PAIRING, HEAD_TO_MANY, + TOPIC_TYPE, MAKE_ID_UNIQUE, KEYWORD_TYPE, INDEX_PAIRING, + DOCUMENT_PUBLISHER_RELATION, PUBLISHER, AFFILIATION_IDENTIFIER_TYPE, COUNTRY_TYPE, + CATEGORY, ACRONYM, RETREIVAL, ATTR_FUNC, ARGS, DOCUMENT_AFFILITATION_RELATION, + AFFILIATION_COUNTRY_RELATION, DOCUMENT_CATEGORY_RELATION, DOCUMENT_ACRONYM_RELATION +) + +# --------------------------- NER default labels +NER_LABELS = [ + "ORG", "PERSON", "GPE", "NORP", "FAC", "LOC", "PRODUCT", + "EVENT", "WORK_OF_ART", "LAW", "LANGUAGE", + "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL", +] + +# ====================================================================================== +# General helpers +# ====================================================================================== +def _map_has_entities(m: Optional[dict]) -> bool: + return bool(m and isinstance(m, dict) and isinstance(m.get("ENTITIES"), list) and len(m["ENTITIES"]) > 0) + +def _run_pass_if_nonempty( + *, + termite, + csv_path: Path, + triplets_path: Path, + triplet_map: dict, + pass_name: str, + verbose: bool = True +) -> bool: + """ + Returns True if the pass ran; False if skipped due to empty ENTITIES. + """ + if not _map_has_entities(triplet_map): + if verbose: + print(f"[TermiteNeo4j] Skipping '{pass_name}' pass: no ENTITIES defined.") + return False + + # Create constraints and run + termite.make_unique_constrains(triplet_map) + termite.from_csv_to_triplets(str(csv_path), str(triplets_path), triplet_map) + termite.update_database_multithreaded(str(triplets_path)) + if verbose: + print(f"[TermiteNeo4j] Finished '{pass_name}' pass -> {triplets_path}") + return True + + +def _get(row, key, default=None): + """Accessor that works for Series/dict/object.""" + try: + if key in row: + return row[key] + except Exception: + pass + if hasattr(row, key): + return getattr(row, key) + if hasattr(row, "get"): + try: + return row.get(key, default) + except Exception: + pass + return default + +def _safe_get(row, key, default=None): + return _get(row, key, default) + +def list_split_no_attrs(data, split_with=';'): + out = [] + if isinstance(data, str): + for v in data.split(split_with): + v = v.strip() + if not v: + continue + e = deepcopy(RETURN_TYPE) + e[ENTITY] = v + out.append(e) + return out or [deepcopy(RETURN_TYPE)] + return [deepcopy(RETURN_TYPE)] + +def split_string(args, split_with=';'): + s = args.get('data') + if isinstance(s, str): + return [t for t in (x.strip() for x in s.split(split_with)) if t] + return [] + +# ====================================================================================== +# Domain extractors (authors, affiliations, categories, acronyms) +# ====================================================================================== + +def make_get_authors_ID_from(ids_col: str, names_col: Optional[str]): + def _fn(args, _ids_col=ids_col, _names_col=names_col): + row = args.get('data', None) + if row is None: + return [deepcopy(RETURN_TYPE)] + data = _safe_get(row, _ids_col, None) + if not isinstance(data, str): + return [deepcopy(RETURN_TYPE)] + ids = [t.strip() for t in data.split(';') if t.strip()] + names = [] + if _names_col: + raw_names = _safe_get(row, _names_col, "") + if isinstance(raw_names, str) and raw_names.strip(): + names = [t.strip() for t in raw_names.split(';')] + out = [] + for i, name in zip(ids, names + [""] * max(0, len(ids) - len(names))): + e = deepcopy(RETURN_TYPE) + e[ENTITY] = i + if name: + e[ATTRIBUTES] = [('name', name)] + out.append(e) + return out or [deepcopy(RETURN_TYPE)] + return _fn + +def _aff_to_dict(cell) -> Dict[str, Dict[str, Any]]: + """Accept dict/list/JSON/string; return dict keyed by affiliation id; ensure name/authors/country keys.""" + if isinstance(cell, (dict, list)): + obj = cell + elif isinstance(cell, str): + s = cell.strip() + if not s or s.lower() == 'nan': + obj = {} + else: + try: + obj = ast.literal_eval(s) + except Exception: + try: + obj = json.loads(s) + except Exception: + obj = {} + else: + obj = {} + + if isinstance(obj, list): + out = {} + for i, item in enumerate(obj): + if not isinstance(item, dict): + continue + key = str(item.get("id", item.get("affiliation_id", i))) + out[key] = { + **item, + "name": item.get("name", item.get("affiliation_name")), + "authors": item.get("authors", item.get("author_ids", [])) or [], + "country": item.get("country", "Unknown"), + } + return out + elif isinstance(obj, dict): + out = {} + for k, val in obj.items(): + if not isinstance(val, dict): + continue + val.setdefault("name", val.get("affiliation_name")) + val.setdefault("authors", val.get("author_ids", [])) + val.setdefault("country", "Unknown") + out[str(k)] = val + return out + return {} + +def make_get_affiliations_from(aff_col: str): + def _fn(args, _aff_col=aff_col): + row = args.get('data', None) + if row is None: + return [deepcopy(RETURN_TYPE)] + raw = _safe_get(row, _aff_col, None) + d = _aff_to_dict(raw) + out = [] + for k, v in d.items(): + e = deepcopy(RETURN_TYPE) + e[ENTITY] = k + e[ATTRIBUTES] = [('name', v.get('name'))] + out.append(e) + return out or [deepcopy(RETURN_TYPE)] + return _fn + +def make_get_countries_from(aff_col: str): + def _fn(args, _aff_col=aff_col): + row = args.get('data', None) + if row is None: + return [deepcopy(RETURN_TYPE)] + raw = _safe_get(row, _aff_col, None) + d = _aff_to_dict(raw) + out = [] + for _, v in d.items(): + e = deepcopy(RETURN_TYPE) + e[ENTITY] = v.get('country', 'Unknown') + out.append(e) + return out or [deepcopy(RETURN_TYPE)] + return _fn + +def get_categories(args): + row = args['data'] + sa = _safe_get(row, 'subject_areas', None) + if isinstance(sa, str): + out = [] + for subj in sa.split(';'): + subj = subj.strip() + if subj: + e = deepcopy(RETURN_TYPE) + e[ENTITY] = subj + out.append(e) + return out or [deepcopy(RETURN_TYPE)] + return [deepcopy(RETURN_TYPE)] + +def get_acronyms(args): + """Tries columns: 'acronym_attribution', 'acronyms', 'acronym'.""" + row = args.get('data', None) + if row is None: + return [] + for col in ('acronym_attribution', 'acronyms', 'acronym'): + v = _safe_get(row, col, None) + if v is not None and str(v).strip() and str(v).lower() != 'nan': + text = str(v).replace(';', ',') + return list_split_no_attrs(text, split_with=',') + return [] + +# ====================================================================================== +# Topics helpers +# ====================================================================================== + +def default_topic_triplet_map(): + return { + 'ENTITIES': [ + {ET: TOPIC_TYPE, MAKE_ID_UNIQUE: True, FROM_COL: 'Graph_Name', + ATTR_COL: [ + {FROM_COL: 'label', ATTR_NAME: 'label'}, + {FROM_COL: 'Graph_Name', ATTR_NAME: 'Graph_Name'}, + ]}, + {ET: KEYWORD_TYPE, MAKE_ID_UNIQUE: True}, + ], + 'RELATIONS': [ + {HT: TOPIC_TYPE, R: 'child_of', TT: TOPIC_TYPE, EXTRACT_T: get_parent_topic}, + {HT: TOPIC_TYPE, R: 'mentions', TT: KEYWORD_TYPE, EXTRACT_T: get_topic_keywords}, + ] + } + +def get_parent_topic(args): + data = args['data'] + graph_name = _get(data, 'Graph_Name') or _get(data, 'graph_name') + if not isinstance(graph_name, str) or not graph_name.strip(): + return [] + gn = graph_name.strip() + parent = gn.rsplit('_', 1)[0] if '_' in gn else None + if not parent: + return [] + e = deepcopy(RETURN_TYPE) + e[ENTITY] = parent + e[ATTRIBUTES] = {"entity_type": "Topic"} + return [e] + +def get_topic_keywords(args): + data = args['data'] + raw = _get(data, 'keywords') or _get(data, 'words') or _get(data, 'keyword_list') + + items: List[str] + if raw is None: + items = [] + elif isinstance(raw, (list, tuple, set)): + items = list(raw) + else: + s = str(raw).strip() + if not s: + items = [] + else: + try: + maybe = ast.literal_eval(s) + if isinstance(maybe, (list, tuple, set)): + items = list(maybe) + else: + items = [maybe] + except Exception: + items = re.split(r'[,\|;]\s*', s) + + seen = set() + keywords = [] + for w in items: + w = str(w).strip() + if w and w not in seen: + seen.add(w) + keywords.append(w) + + out = [] + for w in keywords: + e = deepcopy(RETURN_TYPE) + e[ENTITY] = w + e[ATTRIBUTES] = {"entity_type": "Keyword"} + out.append(e) + return out + +# ====================================================================================== +# NER helpers +# ====================================================================================== + +def _parse_ner_cell(raw): + if raw is None: + return {} + if isinstance(raw, dict): + return raw + s = str(raw).strip() + if not s or s.lower() == "nan": + return {} + try: + return json.loads(s) + except Exception: + try: + v = ast.literal_eval(s) + return v if isinstance(v, dict) else {} + except Exception: + return {} + +def make_ner_extractor(label: str, preferred_cols=None, min_len: int = 1, dedupe: bool = True): + preferred_cols = list(preferred_cols or []) + def _extract(args, _label=label): + row = args.get("data") + if row is None: + return [] + try: + keys = list(getattr(row, "index", [])) or list(getattr(row, "keys", lambda: [])()) + except Exception: + keys = [] + scan_cols = [] + for c in preferred_cols: + if c in keys: + scan_cols.append(c) + for k in keys: + ks = str(k) + if ks == "ner_by_label" or ks.endswith("_ents_by_label") or "ents_by_label" in ks: + if k not in scan_cols: + scan_cols.append(k) + if not scan_cols: + return [] + seen = set() + out = [] + for c in scan_cols: + d = _parse_ner_cell(_safe_get(row, c, None)) + if not isinstance(d, dict): + continue + items = d.get(_label, []) or [] + for text in items: + t = ("" if text is None else str(text)).strip() + if len(t) < min_len: + continue + if dedupe and t in seen: + continue + seen.add(t) + ent = deepcopy(RETURN_TYPE) + ent[ENTITY] = t + ent[ATTRIBUTES] = [("label", _label), ("source_col", str(c))] + out.append(ent) + return out + return _extract + +def default_ner_triplet_map(labels=NER_LABELS, relation_name="mentions", ner_col="ner_by_label", document_id_col="doi"): + m = {"ENTITIES": [], "RELATIONS": []} + m["ENTITIES"].append({ET: DOCUMENT_TYPE, FROM_COL: document_id_col, MAKE_ID_UNIQUE: True}) + for lab in labels: + m["ENTITIES"].append({ET: f"NER_{lab}", MAKE_ID_UNIQUE: True}) + for lab in labels: + m["RELATIONS"].append({ + HT: DOCUMENT_TYPE, R: relation_name, TT: f"NER_{lab}", + EXTRACT_T: make_ner_extractor(lab, preferred_cols=[ner_col], min_len=1, dedupe=True), + PAIRING: HEAD_TO_MANY, + }) + return m + +# ====================================================================================== +# YAML (simple) : head/relation/tail + function +# ====================================================================================== + +def _interpolate_env(value, env): + """Resolve ${path.to.value} inside strings by walking the loaded YAML dict.""" + if isinstance(value, str): + for path in re.findall(r"\$\{([^}]+)\}", value): + cur = env + for p in path.split("."): + cur = cur[p] + value = value.replace("${"+path+"}", str(cur)) + return value + if isinstance(value, list): + return [_interpolate_env(v, env) for v in value] + if isinstance(value, dict): + return {k: _interpolate_env(v, env) for k, v in value.items()} + return value + +def load_settings_yml(path: str) -> dict: + with open(path, "r", encoding="utf-8") as f: + raw = yaml.safe_load(f) or {} + return _interpolate_env(raw, raw) + +def split_factory(sep): + def _split(args, split_with=sep): + return split_string({'data': _safe_get(args, 'data')}, split_with=split_with) + return _split + +def resolve_unpacker(name: str, args_cfg: Optional[dict]) -> Callable: + """ + Supported built-ins: + - split (args: {sep}) + - get_authors_ID (args: {ids_col, names_col}) + - get_affiliations / get_countries (args: {aff_col}) + - get_topic_keywords / get_parent_topic / get_categories / get_acronyms + Optional dynamic loader (commented below) can support "python:module.func". + """ + args_cfg = args_cfg or {} + if name == "split": + return split_factory(args_cfg.get("sep", ";")) + if name == "get_authors_ID": + return make_get_authors_ID_from(args_cfg.get("ids_col", "author_ids"), args_cfg.get("names_col", "authors")) + if name == "get_affiliations": + return make_get_affiliations_from(args_cfg.get("aff_col", "affiliations")) + if name == "get_countries": + return make_get_countries_from(args_cfg.get("aff_col", "affiliations")) + if name == "get_topic_keywords": + return get_topic_keywords + if name == "get_parent_topic": + return get_parent_topic + if name == "get_categories": + return get_categories + if name == "get_acronyms": + return get_acronyms + + # # Optional: dynamic dotted-path loader + # if name.startswith("python:"): + # mod_path = name.split("python:", 1)[1] + # module, func = mod_path.rsplit(".", 1) + # m = __import__(module, fromlist=[func]) + # f = getattr(m, func) + # def _wrapped(args, f=f, a=args_cfg or {}): + # return f(args, **a) if a else f(args) + # return _wrapped + + raise ValueError(f"Unknown function '{name}'") + +def _coalesce_unpacker(spec: Optional[dict] | str) -> Optional[Callable]: + """Accept string ('get_categories') or mapping {function: 'split', args: {...}}.""" + if not spec: + return None + if isinstance(spec, str): + return resolve_unpacker(spec, {}) + fn_name = spec.get("function") or spec.get("fn") + if not fn_name: + return None + return resolve_unpacker(fn_name, spec.get("args", {})) + +def entity_from_simple_yaml(e: dict) -> dict: + """ + Simple keys: + - type (required) -> ET + - from (optional) -> FROM_COL + - unique: bool -> MAKE_ID_UNIQUE + - attrs: list of {from, as, function?} + """ + et = e.get("type") or e.get("et") + if not et: + raise ValueError("Entity requires 'type'") + + out = {ET: et, MAKE_ID_UNIQUE: bool(e.get("unique", e.get("make_id_unique", True)))} + + src = e.get("from", e.get("from_col", None)) + if src is not None: + out[FROM_COL] = src + + attrs = e.get("attrs", []) + if attrs: + cols = [] + for a in attrs: + a_from = a.get("from") or a.get("from_col") + a_as = a.get("as") or a.get("attr_name") + if not a_from or not a_as: + continue + col = {FROM_COL: a_from, ATTR_NAME: a_as} + if a.get("function") or a.get("fn"): + func_spec = a.get("function") or a.get("fn") + func = _coalesce_unpacker(func_spec) + col[RETREIVAL] = func + col[ARGS] = (func_spec.get("args") if isinstance(func_spec, dict) else None) + cols.append(col) + if cols: + out[ATTR_COL] = cols + + if e.get("function") or e.get("fn"): + out[ATTR_FUNC] = _coalesce_unpacker(e.get("function") or e.get("fn")) + if e.get("args"): + out[ARGS] = e["args"] + + return out + +def relation_from_simple_yaml(r: dict) -> dict: + """ + Simple keys: + - head, relation, tail + - head_extract / tail_extract (string or {function, args}) + - pairing (HEAD_TO_MANY, INDEX_PAIRING, etc.) + """ + head = r.get("head") or r.get("ht") + rel = r.get("relation") or r.get("r") + tail = r.get("tail") or r.get("tt") + if not (head and rel and tail): + raise ValueError("Relation requires 'head', 'relation', and 'tail'") + + out = {HT: head, R: rel, TT: tail} + if r.get("pairing"): + out[PAIRING] = r["pairing"] + + he = r.get("head_extract") or r.get("extract_h") + te = r.get("tail_extract") or r.get("extract_t") + if he: + out[EXTRACT_H] = _coalesce_unpacker(he) + if te: + out[EXTRACT_T] = _coalesce_unpacker(te) + return out + +def build_triplet_map_from_simple_yaml(section: dict) -> dict: + """Convert YAML section (data/topics) to Termite triplet-map dict (simple schema).""" + m = {"ENTITIES": [], "RELATIONS": []} + for e in section.get("entities", []): + m["ENTITIES"].append(entity_from_simple_yaml(e)) + for r in section.get("relations", []): + m["RELATIONS"].append(relation_from_simple_yaml(r)) + return m + +def build_ner_triplet_map_from_yaml(ner_cfg: dict, *, doc_col_fallback="doi", ner_col_fallback="ner_by_label") -> dict: + """NER map from YAML; falls back to detected columns if keys omitted.""" + doc_col = ner_cfg.get("document_id_col", doc_col_fallback) + ner_col = ner_cfg.get("ner_col", ner_col_fallback) + relation_name = ner_cfg.get("relation_name", "mentions") + labels = ner_cfg.get("labels", []) + + m = {"ENTITIES": [], "RELATIONS": []} + m["ENTITIES"].append({ET: DOCUMENT_TYPE, FROM_COL: doc_col, MAKE_ID_UNIQUE: True}) + for lab in labels: + m["ENTITIES"].append({ET: f"NER_{lab}", MAKE_ID_UNIQUE: True}) + for lab in labels: + extractor = make_ner_extractor(lab, preferred_cols=[ner_col], min_len=1, dedupe=True) + m["RELATIONS"].append({ + HT: DOCUMENT_TYPE, R: relation_name, TT: f"NER_{lab}", + EXTRACT_T: extractor, PAIRING: HEAD_TO_MANY, + }) + return m + +def merge_maps(base: dict, extra: dict) -> dict: + """Concatenate ENTITIES/RELATIONS lists (Neo4j constraints handle uniqueness).""" + out = {"ENTITIES": list(base.get("ENTITIES", [])), "RELATIONS": list(base.get("RELATIONS", []))} + out["ENTITIES"].extend(extra.get("ENTITIES", [])) + out["RELATIONS"].extend(extra.get("RELATIONS", [])) + return out + + + + +def _prune_empty_node_constraints( + *, + uri: str, + user: str, + password: str, + database: str = "neo4j", + exclude_labels: Optional[Sequence[str]] = None, + verbose: bool = True, +) -> None: + """ + Drops NODE constraints whose target label set currently matches 0 nodes. + Requires Neo4j 4.4+ (SHOW CONSTRAINTS). Uses constraint names for DROP. + """ + try: + from neo4j import GraphDatabase + except Exception as e: + if verbose: + print(f"[TermiteNeo4j] Prune skipped: neo4j driver not installed ({e})") + return + + exclude = set(exclude_labels or []) + + def _bt(s: str) -> str: + # backtick-escape for Cypher identifiers + return s.replace("`", "``") + + driver = GraphDatabase.driver(uri, auth=(user, password)) + try: + with driver.session(database=database) as sess: + rows = sess.run( + "SHOW CONSTRAINTS YIELD name, type, entityType, labelsOrTypes, properties " + "RETURN name, type, entityType, labelsOrTypes, properties" + ).data() + + drop_names: list[str] = [] + for r in rows: + if (r.get("entityType") != "NODE") or not r.get("labelsOrTypes"): + continue + + labels: list[str] = list(r["labelsOrTypes"]) + if any(lab in exclude for lab in labels): + continue + + # Build a label pattern that requires *all* labels (LabelA:LabelB) + label_pattern = ":" + ":".join(f"`{_bt(l)}`" for l in labels) + count_rec = sess.run(f"MATCH (n{label_pattern}) RETURN count(n) AS c").single() + count_val = int(count_rec["c"]) if count_rec else 0 + + if count_val == 0: + name = r.get("name") + if name: + drop_names.append(name) + + for nm in drop_names: + sess.run(f"DROP CONSTRAINT `{_bt(nm)}` IF EXISTS") + if verbose: + print(f"[TermiteNeo4j] Dropped empty-node constraint `{nm}`") + finally: + driver.close() + +# ====================================================================================== +# Mode parsing +# ====================================================================================== + +def _parse_yaml_mode(cfg: dict) -> int: + """ + Returns 1 (only), 2 (merge, default), 3 (ignore). + Accepts numbers or strings: 1/'only', 2/'merge', 3/'ignore'. + """ + v = cfg.get("mode", 2) # default merge + s = str(v).strip().lower() + if s in ("1", "only", "yaml_only", "strict"): + return 1 + if s in ("3", "ignore", "off", "none", "disabled"): + return 3 + return 2 # merge + +# ====================================================================================== +# Block +# ====================================================================================== + +DEFAULT_CALL_SETTINGS: Dict[str, Any] = { + # Source CSVs (None → try bundle keys) + "raw_csv_path": None, # data/docs + "topic_csv_path": None, # topics/labels (NEW) or fallback to data CSV + # Output files + "triplets_filename": "triplets.csv", # legacy (data) + "data_triplets_filename": None, # falls back to triplets_filename + "topic_triplets_filename": "topic_triplets.csv", + "ner_triplets_filename": "ner_triplets.csv", + # Maps + "column_triplet_map": None, + "column_triplet_map_data": None, + "column_triplet_map_topics": None, + "column_triplet_map_ner": None, + # NER config + "ner_labels": NER_LABELS, + "ner_relation_name": "mentions", + "ner_col": "ner_by_label", + # YAML (optional) + "settings_yaml_path": None, + # Neo4j creds + "neo4j_uri": os.getenv("NEO4J_URI", "neo4j://localhost:7666"), + "neo4j_user": os.getenv("NEO4J_USER", "neo4j"), + "neo4j_pass": os.getenv("NEO4J_PASS", "local_password"), + # Optional Termite token + "token": None, + + # --- NEW: constraint pruning --- + "neo4j_db": os.getenv("NEO4J_DB", "neo4j"), # DB name used by SHOW/DROP + "prune_empty_constraints": False, # opt-in + "prune_when": "after", # "before" | "after" | "both" + "prune_labels_exclude": [], # optional: don't drop for these labels +} + +class TermiteNeo4jBlock(AnimalBlock): + """ + Supports YAML `mode`: + 1=only -> use ONLY what’s defined in YAML (no defaults) + 2=merge -> defaults + YAML injections (and run `extra_schemas`) [default] + 3=ignore -> ignore YAML entirely (defaults only; skip extras) + + Builds three base passes (data / topics / NER) and any number of `extra_schemas`. + """ + + CANONICAL_NEEDS: Tuple[str, ...] = ("df", "leaf_labels_csv") + + def __init__( + self, + *, + needs: Sequence[str] = CANONICAL_NEEDS, + provides: Sequence[str] = ("data_triplets_csv", "topic_triplets_csv", "ner_triplets_csv"), + tag: str = "TermiteNeo4j", + init_settings: Optional[Dict[str, Any]] = None, + call_settings: Optional[Dict[str, Any]] = None, + verbose: bool = True, + **kw: Any, + ) -> None: + merged_call_settings = {**DEFAULT_CALL_SETTINGS, **(call_settings or {})} + super().__init__( + needs=needs, + provides=provides, + tag=tag, + init_settings=init_settings or {}, + call_settings=merged_call_settings, + verbose=verbose, + checkpoint=False, + **kw, + ) + + def _prefer_bundle(self, bundle: DataBundle, *keys: str) -> Optional[str]: + for k in keys: + try: + v = bundle.get(k) + except Exception: + v = None + if v: + return v + return None + + def run(self, bundle: DataBundle) -> None: + root_dir = Path(bundle[SAVE_DIR_BUNDLE_KEY]).expanduser().resolve() + out_dir = (root_dir / self.tag).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + + # ---------- Resolve CSV inputs ---------- + data_input = bundle[self.needs[0]] # e.g., a DataFrame (NER-enriched) + out_dir = (Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + + def _ensure_csv_path(obj, fallback_name): + if isinstance(obj, (str, Path)): + return Path(str(obj)).expanduser().resolve() + if hasattr(obj, "to_csv"): + p = out_dir / fallback_name + obj.to_csv(p, index=False, encoding="utf-8-sig") + return p + return Path(str(obj)).expanduser().resolve() + + data_csv_path = _ensure_csv_path(data_input, "termite_input_data.csv") + topic_input = ( + self.call_settings.get("topic_csv_path") + or bundle.get(self.needs[1]) # "leaf_labels_csv" + or data_input + ) + topic_csv_path = _ensure_csv_path(topic_input, "termite_input_topics.csv") + + # ---------- Peek & normalize essential columns ---------- + need_full = False + try: + cols = set(pd.read_csv(data_csv_path, nrows=0).columns) + if not {"doi", "s2id", "eid"} & cols or "year" not in cols: + need_full = True + except Exception: + need_full = True + + df_all = None + if need_full: + df_all = pd.read_csv(data_csv_path) + + if df_all is not None: + if not ({"doi","s2id","eid"} & set(df_all.columns)): + df_all["doc_id"] = [f"row-{i}" for i in range(len(df_all))] + if "year" not in df_all.columns: + df_all["year"] = 0 + df_all.to_csv(data_csv_path, index=False, encoding="utf-8-sig") + cols = set(df_all.columns) + + def _first_present(cands: Sequence[str]) -> Optional[str]: + for c in cands: + if c in cols: + return c + return None + + doc_id_col = _first_present(["doi","s2id","eid"]) or ("doc_id" if "doc_id" in cols else "doi") + + # Build attribute list only from columns that actually exist + attr_cols = [] + if "title" in cols: + attr_cols.append({FROM_COL: "title", ATTR_NAME: "Title"}) + for c, name in (("eid", "EID"), ("s2id", "S2ID"), ("doi", "DOI")): + if c in cols: + attr_cols.append({FROM_COL: c, ATTR_NAME: name}) + + # Detect author id/name columns + author_ids_col = _first_present(["slic_author_ids", "s2_author_ids", "author_ids"]) or "author_ids" + authors_col = _first_present(["slic_authors", "s2_authors", "authors"]) or "authors" + + # Detect affiliations column + aff_col = _first_present(["slic_affiliations", "affiliations"]) or "affiliations" + + # ---------- Output paths ---------- + data_triplets_filename = self.call_settings.get("data_triplets_filename") or self.call_settings.get("triplets_filename", "triplets.csv") + topic_triplets_filename = self.call_settings.get("topic_triplets_filename", "topic_triplets.csv") + ner_triplets_filename = self.call_settings.get("ner_triplets_filename", "ner_triplets.csv") + + data_triplets_path = out_dir / data_triplets_filename + topic_triplets_path = out_dir / topic_triplets_filename + ner_triplets_path = out_dir / ner_triplets_filename + + # ---------- Default triplet maps ---------- + base_data_map = { + 'ENTITIES': [ + {ET: TOPIC_TYPE, MAKE_ID_UNIQUE: True, FROM_COL: 'Graph_Name'}, + {ET: DOCUMENT_TYPE, FROM_COL: doc_id_col, ATTR_COL: attr_cols, MAKE_ID_UNIQUE: True}, + {ET: AFFILIATION_IDENTIFIER_TYPE, MAKE_ID_UNIQUE: True}, + {ET: COUNTRY_TYPE, MAKE_ID_UNIQUE: True}, + {ET: CATEGORY, MAKE_ID_UNIQUE: True}, + {ET: ACRONYM, MAKE_ID_UNIQUE: True}, + {ET: YEAR_TYPE, FROM_COL: 'year', ATTR_COL: None, ATTR_FUNC: None, MAKE_ID_UNIQUE: True}, + { + ET: AUTHOR_ID_TYPE, + FROM_COL: author_ids_col, + ATTR_COL: [{FROM_COL: authors_col, ATTR_NAME: 'Author_Name', RETREIVAL: split_string, ARGS: None}], + ATTR_FUNC: split_string, ARGS: None, MAKE_ID_UNIQUE: True + }, + {ET: PUBLISHER, FROM_COL: 'publication_name', MAKE_ID_UNIQUE: True}, + ], + 'RELATIONS': [ + {HT: DOCUMENT_TYPE, R: 'part_of_topic', TT: TOPIC_TYPE}, + {HT: DOCUMENT_TYPE, R: DOCUMENT_YEAR_RELATION, TT: YEAR_TYPE}, + {HT: AUTHOR_ID_TYPE, R: AUTHOR_DOCUMENT_RELATION, TT: DOCUMENT_TYPE, + EXTRACT_H: make_get_authors_ID_from(author_ids_col, authors_col)}, + {HT: DOCUMENT_TYPE, R: DOCUMENT_AFFILITATION_RELATION, TT: AFFILIATION_IDENTIFIER_TYPE, + EXTRACT_T: make_get_affiliations_from(aff_col)}, + {HT: AFFILIATION_IDENTIFIER_TYPE, R: AFFILIATION_COUNTRY_RELATION, TT: COUNTRY_TYPE, + EXTRACT_H: make_get_affiliations_from(aff_col), + EXTRACT_T: make_get_countries_from(aff_col), PAIRING: INDEX_PAIRING}, + {HT: DOCUMENT_TYPE, R: DOCUMENT_PUBLISHER_RELATION, TT: PUBLISHER}, + {HT: DOCUMENT_TYPE, R: DOCUMENT_CATEGORY_RELATION, TT: CATEGORY, + EXTRACT_T: get_categories, PAIRING: HEAD_TO_MANY}, + {HT: DOCUMENT_TYPE, R: DOCUMENT_ACRONYM_RELATION, TT: ACRONYM, + EXTRACT_T: get_acronyms}, + ], + } + + base_topics_map = default_topic_triplet_map() + base_ner_map = default_ner_triplet_map( + labels=self.call_settings.get("ner_labels", NER_LABELS), + relation_name=self.call_settings.get("ner_relation_name", "mentions"), + ner_col=self.call_settings.get("ner_col", "ner_by_label"), + document_id_col=doc_id_col, + ) + + # ---------- Load YAML (optional) and parse mode ---------- + settings_yaml_path = self.call_settings.get("settings_yaml_path") + yaml_cfg: dict = {} + yaml_root: Optional[Path] = None + yaml_mode: int = 2 # default merge + if settings_yaml_path: + settings_yaml_path = str(Path(settings_yaml_path).expanduser().resolve()) + yaml_root = Path(settings_yaml_path).parent + yaml_cfg = load_settings_yml(settings_yaml_path) or {} + yaml_mode = _parse_yaml_mode(yaml_cfg) + + # optional creds override from YAML (applies for modes 1 and 2) + if yaml_mode in (1, 2) and "neo4j" in yaml_cfg: + self.call_settings["neo4j_uri"] = yaml_cfg["neo4j"].get("uri", self.call_settings["neo4j_uri"]) + self.call_settings["neo4j_user"] = yaml_cfg["neo4j"].get("user", self.call_settings["neo4j_user"]) + self.call_settings["neo4j_pass"] = yaml_cfg["neo4j"].get("pass", self.call_settings["neo4j_pass"]) + else: + yaml_mode = 3 # treat as ignore if no path provided + + # ---------- Build maps per mode ---------- + if yaml_mode == 3: + # IGNORE: behave as if no YAML provided (defaults only; skip extras) + provided_data_map = self.call_settings.get("column_triplet_map_data") or self.call_settings.get("column_triplet_map") + data_triplet_map = provided_data_map or base_data_map + topic_triplet_map = self.call_settings.get("column_triplet_map_topics") or base_topics_map + ner_triplet_map = self.call_settings.get("column_triplet_map_ner") or base_ner_map + extra_schemas_cfg = [] + + elif yaml_mode == 1: + # ONLY YAML: no defaults at all; missing sections -> empty maps + if yaml_cfg.get("data"): + data_triplet_map = build_triplet_map_from_simple_yaml(yaml_cfg["data"]) + else: + data_triplet_map = {"ENTITIES": [], "RELATIONS": []} + + if yaml_cfg.get("topics"): + topic_triplet_map = build_triplet_map_from_simple_yaml(yaml_cfg["topics"]) + else: + topic_triplet_map = {"ENTITIES": [], "RELATIONS": []} + + if yaml_cfg.get("ner"): + ner_triplet_map = build_ner_triplet_map_from_yaml( + yaml_cfg["ner"], doc_col_fallback=doc_id_col, + ner_col_fallback=self.call_settings.get("ner_col", "ner_by_label"), + ) + else: + ner_triplet_map = {"ENTITIES": [], "RELATIONS": []} + + extra_schemas_cfg = yaml_cfg.get("extra_schemas", []) # run exactly as provided + + else: + # MERGE: defaults + YAML injections (and run extras) + if yaml_cfg.get("data"): + yaml_data = build_triplet_map_from_simple_yaml(yaml_cfg["data"]) + data_triplet_map = merge_maps(base_data_map, yaml_data) + else: + provided_data_map = self.call_settings.get("column_triplet_map_data") or self.call_settings.get("column_triplet_map") + data_triplet_map = provided_data_map or base_data_map + + if yaml_cfg.get("topics"): + yaml_topics = build_triplet_map_from_simple_yaml(yaml_cfg["topics"]) + topic_triplet_map = merge_maps(base_topics_map, yaml_topics) + else: + topic_triplet_map = self.call_settings.get("column_triplet_map_topics") or base_topics_map + + if yaml_cfg.get("ner"): + yaml_ner = build_ner_triplet_map_from_yaml( + yaml_cfg["ner"], doc_col_fallback=doc_id_col, + ner_col_fallback=self.call_settings.get("ner_col", "ner_by_label"), + ) + ner_triplet_map = merge_maps(base_ner_map, yaml_ner) + else: + ner_triplet_map = self.call_settings.get("column_triplet_map_ner") or base_ner_map + + extra_schemas_cfg = yaml_cfg.get("extra_schemas", []) + + if yaml_mode in (1, 2): + # allow both under neo4j: {...} and at top-level for convenience + neo4j_cfg = yaml_cfg.get("neo4j", {}) + self.call_settings["neo4j_db"] = neo4j_cfg.get("db", self.call_settings.get("neo4j_db", "neo4j")) + self.call_settings["prune_empty_constraints"] = neo4j_cfg.get( + "prune_empty_constraints", + yaml_cfg.get("prune_empty_constraints", self.call_settings["prune_empty_constraints"]) + ) + self.call_settings["prune_when"] = neo4j_cfg.get( + "prune_when", + yaml_cfg.get("prune_when", self.call_settings.get("prune_when", "after")) + ) + self.call_settings["prune_labels_exclude"] = neo4j_cfg.get( + "prune_labels_exclude", + yaml_cfg.get("prune_labels_exclude", self.call_settings.get("prune_labels_exclude", [])) + ) + + # ---------- Neo4j / Termite ---------- + neo4j_uri = self.call_settings["neo4j_uri"] + neo4j_user = self.call_settings["neo4j_user"] + neo4j_pass = self.call_settings["neo4j_pass"] + token = self.call_settings.get("token", None) + + if self.call_settings.get("prune_empty_constraints") and self.call_settings.get("prune_when") in ("before", "both"): + _prune_empty_node_constraints( + uri=neo4j_uri, + user=neo4j_user, + password=neo4j_pass, + database=self.call_settings.get("neo4j_db", "neo4j"), + exclude_labels=self.call_settings.get("prune_labels_exclude", []), + verbose=self.verbose, + ) + + termite = Termite( + kg_credentials=(neo4j_uri, (neo4j_user, neo4j_pass)), + vector_uri=None, + db_nme="default", + token=token, + verbose=self.verbose, + ) + + # Create uniqueness constraints for base maps up front + # ---------- Base passes (guarded) ---------- + ran_data = _run_pass_if_nonempty( + termite=termite, + csv_path=Path(data_csv_path), + triplets_path=data_triplets_path, + triplet_map=data_triplet_map, + pass_name="data", + verbose=self.verbose, + ) + + ran_topics = _run_pass_if_nonempty( + termite=termite, + csv_path=Path(topic_csv_path), + triplets_path=topic_triplets_path, + triplet_map=topic_triplet_map, + pass_name="topics", + verbose=self.verbose, + ) + + ran_ner = _run_pass_if_nonempty( + termite=termite, + csv_path=Path(data_csv_path), + triplets_path=ner_triplets_path, + triplet_map=ner_triplet_map, + pass_name="ner", + verbose=self.verbose, + ) + + + # ---------- OPTIONAL: any number of extra schema passes (modes 1 & 2 only) ---------- + extra_passes_info: List[Tuple[str, Path]] = [] + if extra_schemas_cfg and isinstance(extra_schemas_cfg, list): + for i, sch in enumerate(extra_schemas_cfg): + if not isinstance(sch, dict): + continue + # Name & CSV + name = str(sch.get("name", f"extra_{i}")).strip() or f"extra_{i}" + csv_arg = sch.get("csv", None) + if csv_arg: + p = Path(csv_arg) + csv_path = (yaml_root / p).resolve() if not p.is_absolute() and yaml_root else p.resolve() + else: + csv_path = data_csv_path # default to main data CSV + + # Build map from the same simple schema keys + has_map_bits = bool(sch.get("entities") or sch.get("relations")) + if has_map_bits: + schema_map = build_triplet_map_from_simple_yaml(sch) + else: + # allow 'map:' nested + if isinstance(sch.get("map"), dict): + schema_map = build_triplet_map_from_simple_yaml(sch["map"]) + else: + continue # nothing to do + + # Output file for this extra pass + extra_triplets_filename = sch.get("triplets_filename", f"{name}_triplets.csv") + extra_triplets_path = out_dir / extra_triplets_filename + + # Execute pass (guarded) + if _run_pass_if_nonempty( + termite=termite, + csv_path=csv_path, + triplets_path=extra_triplets_path, + triplet_map=schema_map, + pass_name=f"extra:{name}", + verbose=self.verbose, + ): + bundle[f"{self.tag}.{name}_triplets_csv"] = extra_triplets_path + extra_passes_info.append((name, extra_triplets_path)) + else: + if self.verbose: + print(f"[{self.tag}] Extra schema '{name}' skipped (no ENTITIES).") + + + # Register & log + bundle[f"{self.tag}.{name}_triplets_csv"] = extra_triplets_path + extra_passes_info.append((name, extra_triplets_path)) + + + if self.call_settings.get("prune_empty_constraints") and self.call_settings.get("prune_when") in ("after", "both"): + _prune_empty_node_constraints( + uri=neo4j_uri, + user=neo4j_user, + password=neo4j_pass, + database=self.call_settings.get("neo4j_db", "neo4j"), + exclude_labels=self.call_settings.get("prune_labels_exclude", []), + verbose=self.verbose, + ) + + # ---------- Register base outputs ---------- + self.register_checkpoint("data_triplets_csv", data_triplets_path) + self.register_checkpoint("topic_triplets_csv", topic_triplets_path) + self.register_checkpoint("ner_triplets_csv", ner_triplets_path) + + bundle[f"{self.tag}.data_triplets_csv"] = data_triplets_path + bundle[f"{self.tag}.topic_triplets_csv"] = topic_triplets_path + bundle[f"{self.tag}.ner_triplets_csv"] = ner_triplets_path + + if self.verbose: + print(f"[{self.tag}] Data triplets @ {data_triplets_path}") + print(f"[{self.tag}] Topic triplets @ {topic_triplets_path}") + print(f"[{self.tag}] NER triplets @ {ner_triplets_path}") + if extra_passes_info: + for nm, pth in extra_passes_info: + print(f"[{self.tag}] Extra schema '{nm}' triplets @ {pth}") diff --git a/TELF/pipeline/blocks/termite_vector_block.py b/TELF/pipeline/blocks/termite_vector_block.py new file mode 100644 index 00000000..dcfd7243 --- /dev/null +++ b/TELF/pipeline/blocks/termite_vector_block.py @@ -0,0 +1,137 @@ +# TELF/pipeline/blocks/termite_vector_index_block.py +from __future__ import annotations +import os +from pathlib import Path +from typing import Any, Dict, Sequence, Tuple, Optional, List + +import pandas as pd +from .base_block import AnimalBlock +from .data_bundle import DataBundle, SAVE_DIR_BUNDLE_KEY +from TELF.applications import Termite + +class TermiteVectorBlock(AnimalBlock): + """ + Mirrors test_termite_e2e.py: + os.environ[...] for OpenSearch + t = Termite(kg_credentials=None, verbose=True, model_name=MODEL) + emb_map = t.compute_embeddings(df, model_name=MODEL) + embeddings = [emb_map[i] for i in df.index] + t.store.ensure_index(index, dim, metric='cosine') + t.store.upsert(index, ids, embeddings, payloads=[{'text': ...}, ...]) + (optional) test search with a query string + + call_settings: + raw_csv_path: str|Path # default: bundle['LeafDataLabels.leaf_data_csv'] + id_column: str # default: 'eid' (falls back to df.index if absent) + text_column: str # default: 'abstract' + index_name: str # default: 'termite_vectors' + model_name: str # default: 'malteos/scincl' + metric: str # default: 'cosine' + env: dict # optional overrides for OS_* and EMBEDDING_STORE + test_query_text: str # optional; if set, do a top-k search + test_k: int # default: 5 + + Provides: + - '{tag}.vector_index_name' + - '{tag}.vector_stats' (docs, dim, metric, index) + - '{tag}.search_hits' (if test_query_text provided) + """ + + CANONICAL_NEEDS: Tuple[str, ...] = ("df",) + + def __init__(self, *, needs: Sequence[str] = CANONICAL_NEEDS, + provides: Sequence[str] = ("vector_index_name", "vector_stats"), + tag: str = "TermiteVectorIndex", + init_settings: Optional[Dict[str, Any]] = None, + call_settings: Optional[Dict[str, Any]] = None, + verbose: bool = True, **kw: Any) -> None: + super().__init__(needs=needs, provides=provides, tag=tag, + init_settings=init_settings or {}, call_settings=call_settings or {}, + verbose=verbose, checkpoint=False, **kw) + + # ---- env like your script ---- + def _init_env(self): + env = { + "EMBEDDING_STORE": "opensearch", + "OS_HOST": os.getenv("OS_HOST", "localhost"), + "OS_PORT": os.getenv("OS_PORT", "9200"), + "OS_USE_SSL": os.getenv("OS_USE_SSL", "false"), + } + overrides = self.call_settings.get("env") or {} + env.update(overrides) + for k, v in env.items(): + os.environ[str(k)] = str(v) + + def _load_df(self, bundle: DataBundle) -> pd.DataFrame: + raw = self.call_settings.get("raw_csv_path") or bundle.get("LeafDataLabels.leaf_data_csv") + if raw: + return pd.read_csv(Path(raw).expanduser().resolve()) + # fallback to bundle df if caller wired it that way + return bundle["df"] + + def run(self, bundle: DataBundle) -> None: + self._init_env() + df = self._load_df(bundle) + + model = self.call_settings.get("model_name", "malteos/scincl") + t = Termite(kg_credentials=None, verbose=self.verbose, model_name=model) + + # embeddings aligned to df.index + emb_map = t.compute_embeddings(df, model_name=model) + embeddings = [emb_map[i] for i in df.index] + if not embeddings: + raise RuntimeError(f"[{self.tag}] No embeddings produced") + dim = len(embeddings[0]) + + index_name = self.call_settings.get("index_name", "termite_vectors") + metric = self.call_settings.get("metric", "cosine") + + # ensure index dimension & metric + t.store.ensure_index(index=index_name, dim=dim, metric=metric) + + # ids + payloads + id_col = self.call_settings.get("id_column", "eid") + ids: List[str] + if id_col in df.columns: + ids = df[id_col].astype(str).tolist() + else: + ids = df.index.astype(str).tolist() # fallback + + text_col = self.call_settings.get("text_column", "abstract") + if text_col not in df.columns: + raise RuntimeError(f"[{self.tag}] text_column '{text_col}' not in DataFrame") + payloads = [{"text": txt} for txt in df[text_col].astype(str).tolist()] + + # upsert + t.store.upsert(index_name, ids, embeddings, payloads=payloads) + + # optional quick search + hits_out = None + qtxt = self.call_settings.get("test_query_text") + if qtxt: + # same embed path as your helper + _qdf = pd.DataFrame({text_col: [qtxt]}) + qmap = t.compute_embeddings(_qdf, model_name=model) + qvec = qmap[_qdf.index[0]] + k = int(self.call_settings.get("test_k", 5)) + hits = t.store.search(index_name, qvec, k=k, source_fields="id,text") + # normalize output for bundle + hits_out = [{"id": _id, "score": float(score), "text": (src or {}).get("text")} for _id, score, src in hits] + bundle[f"{self.tag}.search_hits"] = hits_out + if self.verbose: + print(f"[{self.tag}] Top-{k} search preview:") + for h in hits_out: + print(f" {h['score']:.4f} id={h['id']} text={h['text']}") + + # expose in bundle + bundle[f"{self.tag}.vector_index_name"] = index_name + bundle[f"{self.tag}.vector_stats"] = { + "docs": int(len(df)), + "dim": int(dim), + "metric": metric, + "index": index_name, + "id_column": id_col, + "text_column": text_col, + } + if self.verbose: + print(f"[{self.tag}] Upserted {len(df)} vectors (dim={dim}) into '{index_name}'") diff --git a/TELF/pipeline/blocks/wolf_block.py b/TELF/pipeline/blocks/wolf_block.py index 46af8ee5..085d202c 100644 --- a/TELF/pipeline/blocks/wolf_block.py +++ b/TELF/pipeline/blocks/wolf_block.py @@ -1,10 +1,11 @@ -# wolf_block.py - +# TELF/pipeline/blocks/wolf_block.py from __future__ import annotations from typing import Dict, Sequence, Any, Tuple import os from pathlib import Path +from itertools import combinations +import pandas as pd import numpy as np import networkx as nx from tqdm import tqdm @@ -27,47 +28,47 @@ class WolfBlock(AnimalBlock): """ needs: ['df', 'map'] - provides: ['graph'] + provides: ['graph_'] Automatically checkpoints 'graph' to disk as 'graph.gpickle'. """ - CANONICAL_NEEDS = ('df', 'map') - WOLF_STATS = ['page_rank', 'hubs_authorities', 'betweenness_centrality'] + CANONICAL_NEEDS = ("df", "map") + WOLF_STATS = ["page_rank", "hubs_authorities", "betweenness_centrality"] category_map = { "co-author": { - "col": "slic_author_ids", + "col": "slic_author_ids", "name_col": "slic_authors", - "png": "all_co-authors.png", - "ranks": "co-author_rankings.csv", - "html": "co-authors.html", + "png": "all_co-authors.png", + "ranks": "co-author_rankings.csv", + "html": "co-authors.html", }, "co-affiliation": { - "col": "affiliation_ids", + "col": "affiliation_ids", "name_col": "affiliation_names", - "png": "all_co-affiliations.png", - "ranks": "co-affiliation_rankings.csv", - "html": "co-affiliations.html", + "png": "all_co-affiliations.png", + "ranks": "co-affiliation_rankings.csv", + "html": "co-affiliations.html", }, "co-country": { - "col": "countries", + "col": "countries", "name_col": "countries", - "png": "all_co-countries.png", - "ranks": "co-country_rankings.csv", - "html": "co-countries.html", + "png": "all_co-countries.png", + "ranks": "co-country_rankings.csv", + "html": "co-countries.html", }, } def __init__( self, *, - category: str = 'co-author', + category: str = "co-author", needs: Sequence[str] = CANONICAL_NEEDS, - provides: Sequence[str] = ('graph',), + provides: Sequence[str] = ("graph",), tag: str = "Wolf", conditional_needs: Sequence[Tuple[str, Any]] = (), - init_settings: Dict[str, Any] = None, - call_settings: Dict[str, Any] = None, + init_settings: Dict[str, Any] | None = None, + call_settings: Dict[str, Any] | None = None, verbose: bool = True, ) -> None: if category not in self.category_map: @@ -75,14 +76,12 @@ def __init__( self.category = category - # allow multiple WolfBlock instances without key collision - if provides == ('graph',): + if provides == ("graph",): provides = (f"graph_{self.category}",) default_init = {"verbose": True} - default_call = {} + default_call: Dict[str, Any] = {} - # By default, checkpoint_keys is None → AnimalBlock will use self.provides super().__init__( needs=needs, provides=provides, @@ -93,101 +92,274 @@ def __init__( verbose=verbose, ) + # ----------------------- + # Helpers + # ----------------------- + def _normalize_ids_series(self, s: pd.Series) -> pd.Series: + """ + Normalize an ID column to semicolon-delimited strings. + + - Accept lists/tuples/sets and join with ';' + - Convert common delimiters (',', '|', tab) to ';' + - Trim spaces, drop empties + """ + def _norm(v): + if pd.isna(v): + return "" + if isinstance(v, (list, tuple, set)): + v = ";".join(map(str, v)) + else: + v = str(v) + for sep in [",", "|", "\t"]: + v = v.replace(sep, ";") + parts = [p.strip() for p in v.split(";") if p and p.strip()] + return ";".join(parts) + + return s.map(_norm) + + def _write_empty_artifacts(self, output_dir: Path) -> nx.Graph: + """Create empty outputs (CSV + graph) and return empty graph.""" + stats_path = output_dir / self.category_map[self.category]["ranks"] + pd.DataFrame(columns=["node", *self.WOLF_STATS]).to_csv( + stats_path, index=False, encoding="utf-8-sig" + ) + g = nx.Graph() + graph_path = output_dir / "graph.gpickle" + self.save_path(g, graph_path) + self.register_checkpoint(self.provides[0], graph_path) + return g + + def _build_codep_matrix_fallback(self, s: pd.Series) -> tuple[np.ndarray, list[str]]: + """ + Build a simple symmetric co-occurrence matrix from a normalized + semicolon-delimited ID series. + """ + ids_per_row = [ + [tok for tok in str(v).split(";") if tok] + for v in s.fillna("") + ] + # collect nodes + nodes = [] + seen = set() + for row in ids_per_row: + for tok in row: + if tok not in seen: + seen.add(tok) + nodes.append(tok) + + n = len(nodes) + if n < 2: + return np.zeros((0, 0), dtype=float), [] + + idx = {node: i for i, node in enumerate(nodes)} + X = np.zeros((n, n), dtype=float) + + # count co-occurrences + for row in ids_per_row: + unique_row = sorted(set(row)) + for a, b in combinations(unique_row, 2): + i, j = idx[a], idx[b] + X[i, j] += 1.0 + X[j, i] += 1.0 + + return X, nodes + + def _coerce_node_ids_for_wolf(self, node_ids_any) -> dict | list | tuple | None: + """ + Coerce various node_ids shapes into what Wolf expects: + - dict with sequential-int keys -> OK + - list/tuple of two dicts -> OK (bipartite) + - list/array of labels -> convert to {i: label} + - pandas Index/Series -> convert to {i: label} + - None -> OK + """ + if node_ids_any is None: + return None + + # Already a dict (unipartite) + if isinstance(node_ids_any, dict): + return node_ids_any + + # 2-part structure from some codep implementations + if isinstance(node_ids_any, (list, tuple)) and len(node_ids_any) == 2 \ + and all(isinstance(d, dict) for d in node_ids_any): + return node_ids_any + + # List/array/index/series of labels -> make {i: label} + if isinstance(node_ids_any, (list, tuple, np.ndarray, pd.Index, pd.Series)): + labels = list(node_ids_any) + labels = [str(x) for x in labels] + return {i: lab for i, lab in enumerate(labels)} + + # Fallback: single value? + try: + return {0: str(node_ids_any)} + except Exception: + raise TypeError( + "Unsupported node_ids type; expected dict, (dict, dict), or a sequence of labels." + ) + + # ----------------------- + # Main + # ----------------------- def run(self, bundle: DataBundle) -> None: - # ─── 1) Load the DataFrame & map ────────────────────────────────────── - df = self.load_path(bundle[self.needs[0]]) + # 1) Load inputs + df = self.load_path(bundle[self.needs[0]]) orca_map = bundle[self.needs[1]] - OUTPUT_DIR = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag - output_dir = Path(check_path(os.path.join(OUTPUT_DIR, self.category))) + print("Number of rows in df:", len(df)) + ids_col = self.category_map[self.category]["col"] + if isinstance(df, pd.DataFrame) and ids_col in df.columns: + try: + print("Unique raw values in ID column:", df[ids_col].nunique()) + except Exception: + print("Unique raw values in ID column: (undetermined)") + else: + print(f"Unique IDs: 0 (column '{ids_col}' missing)") - # ─── 2) Build the codependency matrix ──────────────────────────────── - codep = CodependencyMatrixBlock( - col=self.category_map[self.category]['col'] - ) - sub_bundle = DataBundle({ - 'df': df, - SAVE_DIR_BUNDLE_KEY: OUTPUT_DIR - }) - codep(sub_bundle) - X, node_ids = ( - sub_bundle[codep.provides[0]], - sub_bundle[codep.provides[1]] + OUTPUT_ROOT = Path(bundle[SAVE_DIR_BUNDLE_KEY]) / self.tag + output_dir = Path(check_path(os.path.join(OUTPUT_ROOT, self.category))) + output_dir.mkdir(parents=True, exist_ok=True) + + # Guards: 'year' & minimum nodes + if "year" not in df.columns or df["year"].isna().all(): + df = df.copy() + df["year"] = 0 + + # Normalize ID column if present + if ids_col in df.columns: + df = df.copy() + df[ids_col] = self._normalize_ids_series(df[ids_col]) + series = df[ids_col] + else: + # Use a SAFE Series default if the column is missing + series = pd.Series([], dtype=object) + + # Count distinct nodes after normalization + nodes_count = ( + series.dropna() + .astype(str) + .str.split(";") + .explode() + .str.strip() + .replace("", pd.NA) + .dropna() + .nunique() ) + print(f"Usable unique node count after normalization: {nodes_count}") + + if nodes_count < 2: + # produce empty artifacts & exit gracefully + g = self._write_empty_artifacts(output_dir) + bundle[f"{self.tag}.{self.provides[0]}"] = g + print(f"[Wolf/{self.category}] Skipped: only {nodes_count} unique node(s) or column missing.") + return - # ─── 3) Prepare node attributes ────────────────────────────────────── + # 2) Co-dependency matrix + X, node_ids = None, None + try: + codep = CodependencyMatrixBlock( + col=ids_col, + call_settings={ + "split_authors_with": ";", # normalized to ';' + "n_jobs": 1, # robust chunking + }, + ) + sub_bundle = DataBundle({"df": df, SAVE_DIR_BUNDLE_KEY: OUTPUT_ROOT}) + codep(sub_bundle) + + expected = getattr(codep, "provides", ("X", "node_ids")) + if all(k in sub_bundle for k in expected): + X, node_ids = (sub_bundle[expected[0]], sub_bundle[expected[1]]) + else: + print(f"[Wolf/{self.category}] BeaverCodependencyMatrix missing outputs {list(expected)}; using fallback.") + except Exception as e: + print(f"[Wolf/{self.category}] BeaverCodependencyMatrix failed with {type(e).__name__}: {e}") + # fall back below + + # Fallback path if Beaver failed or didn't provide outputs + if X is None or node_ids is None: + X, nodes = self._build_codep_matrix_fallback(series) + if len(nodes) < 2: + g = self._write_empty_artifacts(output_dir) + bundle[f"{self.tag}.{self.provides[0]}"] = g + print(f"[Wolf/{self.category}] Skipped: fallback produced <2 nodes.") + return + node_ids = nodes # list -> will be coerced below + + # ---- COERCE node_ids into the shape Wolf expects ---- + node_ids = self._coerce_node_ids_for_wolf(node_ids) + + # 3) Node attributes wolf = Wolf(**self.init_settings) - wolf.node_ids = node_ids + wolf.node_ids = node_ids # now valid types only if self.category == "co-author": wolf.attributes = create_attributes(orca_map, attribute_names=[]) - - elif self.category == "co-affiliation": - name_col = self.category_map[self.category]['name_col'] - id_col = self.category_map[self.category]['col'] - mapping = get_id_to_name(df, name_col, id_col) - wolf.attributes = {k: {'name': v} for k, v in mapping.items()} - - # ─── 4) Create & annotate the graph ─────────────────────────────────── + name_col = self.category_map[self.category]["name_col"] + id_col = self.category_map[self.category]["col"] + if name_col in df.columns and id_col in df.columns: + mapping = get_id_to_name(df, name_col, id_col) + wolf.attributes = {k: {"name": v} for k, v in mapping.items()} + else: + wolf.attributes = {} + else: + wolf.attributes = {} + + # 4) Create graph & stats graph = wolf.create_graph(X, use_weighted_value=True) for stat in tqdm(self.WOLF_STATS): graph.get_stat(stat) - # ─── 5) Output rankings CSV ────────────────────────────────────────── stats_df = graph.output_stats() - numeric = stats_df.select_dtypes(include=[np.number]).columns - stats_df[numeric] = stats_df[numeric].map(apply_alpha) - stats_df = stats_df \ - .sort_values(by=next(iter(self.WOLF_STATS)), ascending=False) \ - .reset_index(drop=True) + numeric_cols = stats_df.select_dtypes(include=[np.number]).columns + if len(numeric_cols) > 0: + stats_df[numeric_cols] = stats_df[numeric_cols].applymap(apply_alpha) + stats_df = stats_df.sort_values(by=self.WOLF_STATS[0], ascending=False).reset_index(drop=True) stats_df.to_csv( - output_dir / self.category_map[self.category]['ranks'], + output_dir / self.category_map[self.category]["ranks"], index=False, - encoding="utf-8-sig" + encoding="utf-8-sig", ) - # ─── 6) Save full-network plot & HTML ──────────────────────────────── + # 5) Plots graph.visualize( - font_color = 'black', - node_color = '#edede9', - node_size = 100, - highlight_nodes = [], - font_size = 4, - edge_width = 0.08, - figsize = (8, 8), - save_path = str(output_dir / self.category_map[self.category]['png']) + font_color="black", + node_color="#edede9", + node_size=100, + highlight_nodes=[], + font_size=4, + edge_width=0.08, + figsize=(8, 8), + save_path=str(output_dir / self.category_map[self.category]["png"]), ) fig = plot_authors_graph( - df = df, - id_col = self.category_map[self.category]['col'], - name_col = self.category_map[self.category]['name_col'], + df=df, + id_col=self.category_map[self.category]["col"], + name_col=self.category_map[self.category]["name_col"], ) - fig.write_html(str(output_dir / self.category_map[self.category]['html'])) + fig.write_html(str(output_dir / self.category_map[self.category]["html"])) - # ─── 7) Component subplots & word-clouds ───────────────────────────── + # 6) Components & word-clouds save_components( - df = df, - ranking_df = stats_df, - g = graph, - col = self.category_map[self.category]['col'], - results_dir = str(output_dir), + df=df, + ranking_df=stats_df, + g=graph, + col=self.category_map[self.category]["col"], + results_dir=str(output_dir), ) component_wordclouds( - df = df, - g = graph, - col = self.category_map[self.category]['col'], - results_dir = str(output_dir), + df=df, + g=graph, + col=self.category_map[self.category]["col"], + results_dir=str(output_dir), ) - # ─── 8) Checkpoint the graph ───────────────────────────────────────── + # 7) Checkpoint graph_path = output_dir / "graph.gpickle" - graph_path.parent.mkdir(parents=True, exist_ok=True) self.save_path(graph, graph_path) - - # Tell AnimalBlock to record this file under the key "graph" self.register_checkpoint(self.provides[0], graph_path) - # Finally, put the graph into the bundle under your namespaced key - bundle[f"{self.tag}.{self.provides[0]}"] = graph \ No newline at end of file + bundle[f"{self.tag}.{self.provides[0]}"] = graph diff --git a/TELF/post_processing/ArcticFox/arcticfox.py b/TELF/post_processing/ArcticFox/arcticfox.py index 378c7c6f..abbd37e7 100644 --- a/TELF/post_processing/ArcticFox/arcticfox.py +++ b/TELF/post_processing/ArcticFox/arcticfox.py @@ -1,9 +1,13 @@ -from .helpers.local_labels import ClusterLabeler +from .helpers.local_labels import ClusterLabeler from .helpers.intial_post_process import HNMFkPostProcessor from .helpers.post_statistics import HNMFkStatsGenerator from pathlib import Path +from typing import Optional, Sequence, Literal, Set import pandas as pd +Step = Literal["post", "label", "stats"] + + class ArcticFox: def __init__( self, @@ -51,30 +55,95 @@ def run_full_pipeline( self, vocab, data_df, - text_column=None, - ollama_model="llama3.2:3b-instruct-fp16", - label_clusters=True, - generate_stats=True, - generate_visuals=True, - process_parents=True, - skip_completed=True, + text_column: Optional[str] = None, + ollama_model: str = "llama3.2:3b-instruct-fp16", + label_clusters: bool = True, + generate_stats: bool = True, + generate_visuals: bool = True, # kept for backwards-compatibility + process_parents: bool = True, + skip_completed: bool = True, label_criteria=None, label_info=None, - number_of_labels=5 + number_of_labels: int = 5, + # NEW: choose exact subset of steps to run; None keeps legacy behavior + steps: Optional[Sequence[Step]] = None, ): - text_column = text_column or self.clean_cols_name + """ + Run any subset of the pipeline while preserving order: - print("Step 1: Post-processing W/H matrix and cluster data...") - self.postprocessor.post_process_hnmfk( - hnmfk_model=self.model, - V=vocab, - D=data_df, - col_name=text_column, - skip_completed=skip_completed, - process_parents=process_parents - ) + 'post' → post_process_hnmfk + 'label' → _label_all_clusters (requires 'post' artifacts) + 'stats' → generate_cluster_stats (requires 'post' artifacts) - if label_clusters: + Rules: + • 'label' and/or 'stats' can be run without 'post' only if artifacts already exist. + • Order is always post → label → stats, even if you request multiple. + """ + text_column = text_column or self.clean_cols_name + + # ---- resolve which steps to run, validate names ---- + if steps is not None: + steps_set: Set[Step] = set(steps) + invalid = steps_set.difference({"post", "label", "stats"}) + if invalid: + raise ValueError(f"Invalid steps: {sorted(invalid)}; allowed: 'post','label','stats'") + do_post = "post" in steps_set + do_label = "label" in steps_set + do_stats = "stats" in steps_set + else: + # Back-compat defaults using the existing booleans + do_post = True + do_label = bool(label_clusters) + do_stats = bool(generate_stats) + + # ---- helper: ensure post-processing outputs exist if required ---- + def _assert_postprocessed_ready() -> None: + missing = [] + for node in self.model.traverse_nodes(): + if node["leaf"] or process_parents: + w = node.get('W') + if w is None: + sig = node.get('signature') + if sig is None: + missing.append(f"{node.get('node_save_path', '')} (no W or signature)") + continue + w = sig.reshape(-1, 1) + + k = w.shape[1] + node_dir = Path(node["node_save_path"]).resolve().parent + cluster_file = node_dir / f"cluster_for_k={k}.csv" + top_words_file = node_dir / "top_words.csv" + if not (cluster_file.exists() and top_words_file.exists()): + missing.append(str(node_dir)) + + if missing: + raise RuntimeError( + "Labeling and/or stats require post-processing artifacts that were not found:\n" + + "\n".join(f" - {m}" for m in missing) + + "\nInclude 'post' in `steps`, or run the post-processing step first." + ) + + # ───────────────────────────────────────────────────────────── + # Step 1: POST-PROCESS (optional) + # ───────────────────────────────────────────────────────────── + if do_post: + print("Step 1: Post-processing W/H matrix and cluster data...") + self.postprocessor.post_process_hnmfk( + hnmfk_model=self.model, + V=vocab, + D=data_df, + col_name=text_column, + skip_completed=skip_completed, + process_parents=process_parents + ) + else: + if do_label or do_stats: + _assert_postprocessed_ready() + + # ───────────────────────────────────────────────────────────── + # Step 2: LABEL (optional; never before post) + # ───────────────────────────────────────────────────────────── + if do_label: print("Step 2: Labeling clusters with LLM...") self._label_all_clusters( vocab=vocab, @@ -87,7 +156,10 @@ def run_full_pipeline( process_parents=process_parents ) - if generate_stats: + # ───────────────────────────────────────────────────────────── + # Step 3: STATS (optional; always last) + # ───────────────────────────────────────────────────────────── + if do_stats: print("Step 3: Generating Peacock visual stats...") self.stats_generator.generate_cluster_stats( model=self.model, @@ -95,16 +167,18 @@ def run_full_pipeline( skip_completed=skip_completed ) - def _label_all_clusters( self, vocab, data_df, text_column, ollama_model, label_criteria, label_info, number_of_labels, process_parents ): for node in self.model.traverse_nodes(): if node["leaf"] or process_parents: - w = node['W'] + w = node.get('W') if w is None: - w = node['signature'].reshape(-1, 1) + sig = node.get('signature') + if sig is None: + continue + w = sig.reshape(-1, 1) node_dir = Path(node["node_save_path"]).resolve().parent cluster_file = node_dir / f"cluster_for_k={w.shape[1]}.csv" @@ -132,6 +206,7 @@ def _label_all_clusters( for k, v in annotations.items() ]).to_csv(node_dir / "cluster_summaries.csv", index=False) + # Convenience one-offs (unchanged) def run_labeling(self, df, top_words_df, ollama_model_name, label_criteria=None, additional_info=None, number_of_labels=5): return self.labeler.label_clusters_ollama( top_words_df=top_words_df, diff --git a/TELF/pre_processing/Orca/orca.py b/TELF/pre_processing/Orca/orca.py index 43bb3b53..682e5302 100644 --- a/TELF/pre_processing/Orca/orca.py +++ b/TELF/pre_processing/Orca/orca.py @@ -1,394 +1,447 @@ +from __future__ import annotations +# TELF/pre_processing/Orca/orca.py import os import ast import copy import pickle -import warnings +import warnings import pandas as pd import networkx as nx -from tqdm import tqdm +from tqdm import tqdm +# Minimal no-op AuthorMatcher fallback +# Path 1 (temporary to match your current import resolution): +# TELF/pipeline/blocks/AuthorMatcher.py +# Path 2 (correct long-term location): +# TELF/pre_processing/Orca/AuthorMatcher.py -from .AuthorMatcher import AuthorMatcher +import pandas as pd +from typing import Dict, Iterable, List, Tuple, Union -class Orca: +class AuthorMatcher: + """ + Fallback stub used when the real AuthorMatcher isn't available. + Produces an empty matches DataFrame (or builds rows from known_matches if you provide them). + This is enough for Orca to proceed via the Scopus-only/S2-only residual logic. + """ + + def __init__(self, df: pd.DataFrame, n_jobs: int = -1, verbose: bool = False): + self.df = df + self.n_jobs = n_jobs + self.verbose = verbose - #Code where we have a precomputed duplicates. - # Pre-computed Scopus duplicates file containing entries. - # If no duplicates are computed with DAF, this file will be used instead for duplicate removal - #DUPLICATES_1M = 'scopus_1m_cited_collab_matches.p' + def match(self, known_matches: Dict[str, Union[str, Iterable[str]]] | None = None) -> pd.DataFrame: + cols = ["S2_Author_ID", "S2_Author_Name", "SCOPUS_Author_ID"] + if not known_matches: + # No matches known → return empty, Orca will handle residual mapping. + return pd.DataFrame(columns=cols) + + # Build a quick s2_id -> name map if available + s2_name_map: Dict[str, str] = {} + try: + if {"s2_author_ids", "s2_authors"}.issubset(self.df.columns): + tmp = self.df[["s2_author_ids", "s2_authors"]].dropna(how="any") + for ids, names in zip(tmp["s2_author_ids"], tmp["s2_authors"]): + ids = str(ids).split(";") + names = str(names).split(";") + for i, n in zip(ids, names): + s2_name_map.setdefault(i, n) + except Exception: + pass + + rows: List[Dict[str, str]] = [] + for k, vs in known_matches.items(): + if not isinstance(vs, (list, tuple, set)): + vs = [vs] + for v in vs: + k_str, v_str = str(k), str(v) + + # Heuristics to assign which side is S2 vs Scopus (best effort) + k_in_s2 = k_str in s2_name_map + v_in_s2 = v_str in s2_name_map + if k_in_s2 and not v_in_s2: + s2_id, scopus_id = k_str, v_str + elif v_in_s2 and not k_in_s2: + s2_id, scopus_id = v_str, k_str + else: + # Ambiguous → skip quietly + continue + + rows.append( + { + "S2_Author_ID": s2_id, + "S2_Author_Name": s2_name_map.get(s2_id, "Unknown"), + "SCOPUS_Author_ID": scopus_id, + } + ) + + return pd.DataFrame(rows, columns=cols) + + +class Orca: + """ + Construct SLIC author ids + apply them to a SLIC-style paper dataframe. + """ def __init__(self, duplicates=None, s2_duplicates=None, verbose=False): self.slic_df = None self.duplicates = duplicates self.s2_duplicates = s2_duplicates self.verbose = verbose - - - def _run_scopus(self, df): - """ - Helper function for creating a SLIC map file for a dataset that only contains Scopus information - """ - # generate a map of scopus ids to affiliations - affiliations_map = self.__generate_affiliations_map(df) - - # generate author maps - scopus_author_map = self.__generate_author_map(df, 'author_ids', 'authors') - - # correct duplicates - duplicates = {} - for entry in self.duplicates: - if not entry & scopus_author_map.keys(): - continue - - entry = entry.copy() # avoid modifying duplicates in place - best_id = sorted(entry, key=lambda x: len(scopus_author_map.get(x, '')), reverse=True)[0] - merged_affiliations = self.__merge_scopus_affiliations(entry, affiliations_map) - affiliations_map[best_id] = merged_affiliations - # remove old duplicate ids from both maps - entry.remove(best_id) - for x in entry: - del scopus_author_map[x] - del affiliations_map[x] - - # add duplicates for tracking purposes - duplicates[best_id] = list(entry) - - - ## generate SLIC IDs - slic_count = 0 - slic_df = { - 'slic_id': [], - 'slic_name': [], - 'scopus_ids': [], - 'scopus_names': [], - 'scopus_affiliations': [], - 's2_ids': [], - 's2_names': [], - } - - for i, scopus_id in enumerate(scopus_author_map): - scopus_name = scopus_author_map.get(scopus_id) - scopus_affiliations = affiliations_map.get(scopus_id) - if scopus_id in duplicates: - scopus_id = ';'.join([scopus_id] + duplicates[scopus_id]) - - slic_df['slic_id'].append(f'S{i}') - slic_df['slic_name'].append(scopus_name) - slic_df['scopus_ids'].append(scopus_id) - slic_df['scopus_names'].append(scopus_name) - slic_df['scopus_affiliations'].append(scopus_affiliations) - slic_df['s2_ids'].append(None) - slic_df['s2_names'].append(None) - - slic_df = pd.DataFrame.from_dict(slic_df) - return slic_df - - + # ───────────────────────────────────────────────────────────────────────── + # Public API + # ───────────────────────────────────────────────────────────────────────── + def run(self, df, scopus_duplicates=None, s2_duplicates=None, known_matches=None, n_jobs=-1): """ - Run Orca and form SLIC ids for a given dataset - - Parameters - ---------- - df: pandas.DataFrame - The SLIC dataframe for which author SLIC ids need to be created - scopus_duplicates: list(set), optional - A list of sets where each set contains scopus author ids that refer to the same person. In the ideal case, each - author only has one scopus id. However, this ideal does not hold up in practice and some authors are represented - by two or more scopus ids. Duplicate authors can be found using the Orca.DuplicateAuthorFinder tool. If not provided, - a pre-computed scopus duplicate map is used (pre-computed from 1 million Scopus papers). If provided, this map is - overriden by the user input. Default is None. - s2_duplicates: list(set), optional - A list of sets where each set contains s2 author ids that refer to the same person. If not provided, s2 author ids - are not scanned for duplicates / only compared against scopus matches as duplicate detection. Default is None. - known_matches: dict, optional - A dict of s2 id keys to scopus id values. This dictionary is used to override the author matching if groundtruth - is known. This is useful for helping the tool work around edge cases. Default is None. - - Returns - ------- - None + Form the SLIC map from Scopus-only, S2-only, or hybrid dataframes. """ - # process duplicates (if passed) if scopus_duplicates is not None: self.duplicates = scopus_duplicates if s2_duplicates is not None: self.s2_duplicates = s2_duplicates - - valid, error = self.__verify_df(df) # make sure that the passed dataframe meets expected - if not valid: - if False: # TODO: set valid flag to check if Scopus only data here - raise ValueError(error) - else: - self.slic_df = self._run_scopus(df) - return self.slic_df - - # generate a map of scopus ids to affiliations + + has_scopus = {"author_ids", "authors"}.issubset(df.columns) + has_s2 = {"s2_author_ids", "s2_authors"}.issubset(df.columns) + + if has_scopus and not has_s2: + self.slic_df = self._run_scopus(df) + return self.slic_df + if has_s2 and not has_scopus: + self.slic_df = self._run_s2_only(df) + return self.slic_df + if not (has_scopus or has_s2): + raise ValueError( + "Orca.run(): DataFrame must contain Scopus ('author_ids','authors') " + "or S2 ('s2_author_ids','s2_authors') columns." + ) + + # Hybrid path affiliations_map = self.__generate_affiliations_map(df) + s2_author_map = self.__generate_author_map(df, "s2_author_ids", "s2_authors") + scopus_author_map = self.__generate_author_map(df, "author_ids", "authors") - # generate author maps - s2_author_map = self.__generate_author_map(df, 's2_author_ids', 's2_authors') - scopus_author_map = self.__generate_author_map(df, 'author_ids', 'authors') - - # match scopus author ids to s2 author ids known_matches = {} if not known_matches else known_matches am = AuthorMatcher(df, n_jobs=n_jobs, verbose=self.verbose) am_df = am.match(known_matches=known_matches) - - # process scopus duplicates + + # Enrich with known Scopus duplicates am_enriched = self.__add_scopus_duplicates(am_df, self.duplicates) am_df = pd.concat([am_df, am_enriched], axis=0, ignore_index=True) - ## generate SLIC IDs + # Build components across S2/Scopus ids slic_count = 0 + seen_s2, seen_scopus = set(), set() + matches = self.__uncouple_author_matches(am_df, self.s2_duplicates) + slic_df = { - 'slic_id': [], - 'slic_name': [], - 'scopus_ids': [], - 'scopus_names': [], - 'scopus_affiliations': [], - 's2_ids': [], - 's2_names': [], + "slic_id": [], + "slic_name": [], + "scopus_ids": [], + "scopus_names": [], + "scopus_affiliations": [], + "s2_ids": [], + "s2_names": [], } - - # 1. assign SLIC IDs to author ids that have correspondence between s2 and scopus - seen_s2 = set() - seen_scopus = set() - matches = self.__uncouple_author_matches(am_df, self.s2_duplicates) + + # 1) Matched S2<->Scopus groups for entry in matches: - s2_ids = entry['s2'] - s2_names = {s2_author_map.get(x, 'Unknown') for x in s2_ids if x in s2_author_map} - scopus_ids = entry['scopus'] - scopus_names = {scopus_author_map[x] for x in scopus_ids if x in scopus_author_map} - scopus_affiliations = self.__merge_scopus_affiliations(scopus_ids, affiliations_map) - slic_name = entry['name'] - - seen_s2 |= s2_ids - seen_scopus |= scopus_ids - - slic_df['slic_id'].append(f'S{slic_count}') - slic_df['slic_name'].append(slic_name) - slic_df['scopus_ids'].append(';'.join(scopus_ids)) - slic_df['scopus_names'].append(';'.join(scopus_names)) - slic_df['scopus_affiliations'].append(scopus_affiliations) - slic_df['s2_ids'].append(';'.join(s2_ids)) - slic_df['s2_names'].append(';'.join(s2_names)) + s2_id_set = entry["s2"] + s2_name_set = {s2_author_map.get(x, "Unknown") for x in s2_id_set if x in s2_author_map} + scopus_id_set = entry["scopus"] + scopus_name_set = {scopus_author_map[x] for x in scopus_id_set if x in scopus_author_map} + scopus_affiliations = self.__merge_scopus_affiliations(scopus_id_set, affiliations_map) + slic_name = entry["name"] + + seen_s2 |= s2_id_set + seen_scopus |= scopus_id_set + + slic_df["slic_id"].append(f"S{slic_count}") + slic_df["slic_name"].append(slic_name) + slic_df["scopus_ids"].append(";".join(sorted(scopus_id_set))) + slic_df["scopus_names"].append(";".join(sorted(scopus_name_set))) + slic_df["scopus_affiliations"].append(scopus_affiliations) + slic_df["s2_ids"].append(";".join(sorted(s2_id_set))) + slic_df["s2_names"].append(";".join(sorted(s2_name_set))) slic_count += 1 - - # 2. assign SLIC IDs to scopus ids that did not have correspondence - df_scopus_authors = {x for y in df.author_ids.to_list() if not pd.isna(y) for x in y.split(';')} - df_scopus_authors -= seen_scopus - for scopus_id in df_scopus_authors: - scopus_name = scopus_author_map.get(scopus_id, None) - - slic_df['slic_id'].append(f'S{slic_count}') - slic_df['slic_name'].append(scopus_name) - slic_df['scopus_ids'].append(scopus_id) - slic_df['scopus_names'].append(scopus_name) - slic_df['scopus_affiliations'].append(affiliations_map.get(scopus_id, None)) - slic_df['s2_ids'].append(None) - slic_df['s2_names'].append(None) - slic_count += 1 - - # 3. assign SLIC Ids to remaining s2 authors - df_s2_authors = {x for y in df.s2_author_ids.to_list() if not pd.isna(y) for x in y.split(';')} - df_s2_authors -= seen_s2 - for s2_id in df_s2_authors: - if s2_id in seen_s2: - continue - - # update the common fields - slic_df['slic_id'].append(f'S{slic_count}') - slic_df['scopus_ids'].append(None) - slic_df['scopus_names'].append(None) - slic_df['scopus_affiliations'].append(None) - - # handle s2 duplicates if they exist - s2_dup_ids = self.s2_duplicates.get(s2_id) - if s2_dup_ids is not None: - s2_ids = {s2_id} | set(s2_dup_ids) - s2_names = {s2_author_map.get(x, 'Unknown') for x in s2_ids if x in s2_author_map} - if s2_names == {'Unknown'}: # handle case where all name missing - s2_names = set() - - # get the slic name - try: - slic_name = max(s2_names, key=len) - slic_name = None if slic_name == 'Unknown' else slic_name - except ValueError: - slic_name = None - - # update the data map - slic_df['slic_name'].append(slic_name) - s2_ids_str = ';'.join(s2_ids) if s2_ids else None - slic_df['s2_ids'].append(s2_ids_str) - s2_names = ';'.join(s2_names) if s2_names else None - slic_df['s2_names'].append(s2_names) - seen_s2 |= s2_ids - - else: - s2_name = s2_author_map.get(s2_id, None) - slic_df['slic_name'].append(s2_name) - slic_df['s2_ids'].append(s2_id) - slic_df['s2_names'].append(s2_name) - seen_s2 |= s2_ids - - # incremenet the slic id identifier - slic_count += 1 - + + # 2) Scopus-only residuals + if "author_ids" in df.columns: + df_scopus_authors = {x for y in df.author_ids.to_list() if not pd.isna(y) for x in y.split(";")} + df_scopus_authors -= seen_scopus + for scopus_id in sorted(df_scopus_authors): + scopus_name = scopus_author_map.get(scopus_id, None) + slic_df["slic_id"].append(f"S{slic_count}") + slic_df["slic_name"].append(scopus_name) + slic_df["scopus_ids"].append(scopus_id) + slic_df["scopus_names"].append(scopus_name) + slic_df["scopus_affiliations"].append(affiliations_map.get(scopus_id, None)) + slic_df["s2_ids"].append(None) + slic_df["s2_names"].append(None) + slic_count += 1 + + # 3) S2-only residuals + if "s2_author_ids" in df.columns: + df_s2_authors = {x for y in df.s2_author_ids.to_list() if not pd.isna(y) for x in y.split(";")} + df_s2_authors -= seen_s2 + for s2_id in sorted(df_s2_authors): + slic_df["slic_id"].append(f"S{slic_count}") + slic_df["scopus_ids"].append(None) + slic_df["scopus_names"].append(None) + slic_df["scopus_affiliations"].append(None) + + s2_dup_ids = self.s2_duplicates.get(s2_id) + if s2_dup_ids is not None: + s2_ids_all = {s2_id} | set(s2_dup_ids) + s2_author_map_local = {x: s2_author_map.get(x, "Unknown") for x in s2_ids_all} + s2_name_set = {v for v in s2_author_map_local.values() if v != "Unknown"} + slic_name = max(s2_name_set, key=len) if s2_name_set else None + + slic_df["slic_name"].append(slic_name) + slic_df["s2_ids"].append(";".join(sorted(s2_ids_all)) if s2_ids_all else None) + slic_df["s2_names"].append(";".join(sorted(s2_name_set)) if s2_name_set else None) + seen_s2 |= s2_ids_all + else: + s2_name = s2_author_map.get(s2_id, None) + slic_df["slic_name"].append(s2_name) + slic_df["s2_ids"].append(s2_id) + slic_df["s2_names"].append(s2_name) + seen_s2.add(s2_id) + + slic_count += 1 + slic_df = pd.DataFrame.from_dict(slic_df) - slic_df = slic_df.loc[slic_df.slic_name != 'Unknown'].copy().reset_index(drop=True) + slic_df = slic_df.loc[slic_df.slic_name != "Unknown"].copy().reset_index(drop=True) self.slic_df = slic_df return slic_df - - + def apply(self, df, slic_df=None): """ - Apply the SLIC id mapping to a SLIC papers dataframe - - Parameters - ---------- - df: pandas.DataFrame - The SLIC dataframe for which author SLIC ids need to be created - slic_df: pandas.DataFrame, optional - A pre-computed DataFrame with SLIC id mappings. This parameter is provided in the rare cases that a SLIC map is - being used between multiple datasets (i.e. dataset B is a subset of A and slic_df was computed for A). Be aware that - setting a value for slic_df is not recommended! If using this parameter, verify that all desired scopus/s2 authors - have existing SLIC ids. To be sure of the validity of your results, use Orca.run() before using Orca.apply() and - do not pass a value for this parameter. - - Returns - ------- - orca_df: pandas.DataFrame - df with standarized author information (columns for 'SLIC_ids' and 'SLIC_affiliations') + Apply the SLIC id mapping to a SLIC papers dataframe. + Keeps papers even when SLIC author ids are missing (warns only). """ if slic_df is None and self.slic_df is None: - return ValueError('No SLIC ID map found. First, compute the map with Orca.run()') + return ValueError("No SLIC ID map found. First, compute the map with Orca.run()") if slic_df is not None and self.slic_df is not None: - warnings.warn('[Orca]: slic_df was passed as an argument however this Orca object already has a ' \ - 'stored slic_df object.\n\t\tOverwriting stored slic_df with given argument. If this ' \ - 'message is unexpected, use Orca.apply() without specifying `slic_df`', RuntimeWarning) - + warnings.warn( + "[Orca]: slic_df was passed as an argument however this Orca object already has a " + "stored slic_df object.\n\t\tOverwriting stored slic_df with given argument. If this " + "message is unexpected, use Orca.apply() without specifying `slic_df`", + RuntimeWarning, + ) + if slic_df is not None: self.slic_df = slic_df - - # verify that paper scopus ids and s2ids are unique - if 'eid' in df.columns and df.eid.nunique() != len(df.loc[~df.eid.isnull()]): - df = df[~df['eid'].duplicated(keep='first') | df['eid'].isna()].copy() - warnings.warn('[Orca]: Encountered duplicate Scopus IDs (`eid`) in df. Dropping duplicate papers.') - if 's2id' in df.columns and df.s2id.nunique() != len(df.loc[~df.s2id.isnull()]): - df = df[~df['s2id'].duplicated(keep='first') | df['s2id'].isna()].copy() - warnings.warn('[Orca]: Encountered duplicate S2 IDs (`s2id`) in df. Dropping duplicate papers.') - - # replace scopus and s2 author ids respectively - if 's2id' in df.columns and 'eid' in df.columns: + + # Uniqueness guards + if "eid" in df.columns and df.eid.nunique() != len(df.loc[~df.eid.isnull()]): + df = df[~df["eid"].duplicated(keep="first") | df["eid"].isna()].copy() + warnings.warn("[Orca]: Encountered duplicate Scopus IDs (`eid`) in df. Dropping duplicate papers.") + if "s2id" in df.columns and df.s2id.nunique() != len(df.loc[~df.s2id.isnull()]): + df = df[~df["s2id"].duplicated(keep="first") | df["s2id"].isna()].copy() + warnings.warn("[Orca]: Encountered duplicate S2 IDs (`s2id`) in df. Dropping duplicate papers.") + + # Compute per-source SLIC ids + if "s2id" in df.columns and "eid" in df.columns: scopus_df = self.__compute_slic_scopus(df) s2_df = self.__compute_slic_s2(df) - - # merge and build output dataframe - df2 = pd.merge(df, scopus_df, on='eid', how='outer') - df3 = pd.merge(df2, s2_df, on='s2id', how='outer') + df2 = pd.merge(df, scopus_df, on="eid", how="outer") + df3 = pd.merge(df2, s2_df, on="s2id", how="outer") orca_df = df3.copy() - orca_df['slic_author_ids'] = orca_df['slic_author_ids_x'].combine_first(orca_df['slic_author_ids_y']) - orca_df = orca_df.drop(columns=['slic_author_ids_x', 'slic_author_ids_y']) - - elif 's2id' in df.columns: + + # unify slic_author_ids + orca_df["slic_author_ids"] = orca_df["slic_author_ids_x"].combine_first(orca_df["slic_author_ids_y"]) + orca_df = orca_df.drop(columns=["slic_author_ids_x", "slic_author_ids_y"]) + + # unify slic_affiliations if present as suffixes + if "slic_affiliations_x" in orca_df.columns or "slic_affiliations_y" in orca_df.columns: + left = orca_df.get("slic_affiliations_x") + right = orca_df.get("slic_affiliations_y") + if left is not None and right is not None: + orca_df["slic_affiliations"] = left.combine_first(right) + orca_df.drop( + columns=[c for c in ["slic_affiliations_x", "slic_affiliations_y"] if c in orca_df.columns], + inplace=True, + ) + elif left is not None: + orca_df.rename(columns={"slic_affiliations_x": "slic_affiliations"}, inplace=True) + else: + orca_df.rename(columns={"slic_affiliations_y": "slic_affiliations"}, inplace=True) + + elif "s2id" in df.columns: s2_df = self.__compute_slic_s2(df) - orca_df = pd.merge(df, s2_df, on='s2id', how='outer') + orca_df = pd.merge(df, s2_df, on="s2id", how="outer") + if "slic_affiliations" not in orca_df.columns: + orca_df["slic_affiliations"] = None + else: scopus_df = self.__compute_slic_scopus(df) - orca_df = pd.merge(df, scopus_df, on='eid', how='outer') - - if orca_df.slic_author_ids.isna().any(): - original_len = len(orca_df) - orca_df.dropna(subset=['slic_author_ids'], inplace=True) - warnings.warn(f'[Orca]: Found {original_len - len(orca_df)} papers with missing ' \ - 'SLIC author IDs. Dropping these papers.') - - # add a column of slic author names using the matched slic ids - slic_authors = {k:v for k,v in zip(self.slic_df.slic_id.to_list(), self.slic_df.slic_name.to_list())} + orca_df = pd.merge(df, scopus_df, on="eid", how="outer") + if "slic_affiliations_x" in orca_df.columns: + orca_df.rename(columns={"slic_affiliations_x": "slic_affiliations"}, inplace=True) + if "slic_affiliations_y" in orca_df.columns: + orca_df.drop(columns=["slic_affiliations_y"], inplace=True) + + # Keep papers with missing SLIC ids (warn only) + if "slic_author_ids" in orca_df.columns: + missing = int(orca_df["slic_author_ids"].isna().sum()) + if missing: + warnings.warn(f"[Orca]: {missing} papers have no SLIC author IDs (S2-only or unmatched). Keeping them.") + + if "slic_affiliations" not in orca_df.columns: + orca_df["slic_affiliations"] = None + + # Add SLIC author names + slic_authors = {k: v for k, v in zip(self.slic_df.slic_id.to_list(), self.slic_df.slic_name.to_list())} + def map_ids_to_names(ids): if pd.isna(ids): return None - names = [slic_authors.get(str(i), '') for i in ids.split(';')] + names = [slic_authors.get(str(i), "") for i in ids.split(";")] names = [name for name in names if name] - return ';'.join(names) - orca_df['slic_authors'] = orca_df['slic_author_ids'].apply(map_ids_to_names) + return ";".join(names) + + orca_df["slic_authors"] = orca_df["slic_author_ids"].apply(map_ids_to_names) return orca_df.reset_index(drop=True) - - - def __verify_df(self, df): - """ - Verify that the given papers dataframe matches the SLIC standard and can be used with Orca - - Parameters - ---------- - df: pandas.DataFrame - The SLIC papers DataFrame for which author SLIC ids need to be created - - Returns - ------- - flag: bool - If true, df passes the test and can be used with Orca - error: str, None - If flag is True, None is returned. Otherwise a string with the encountered error is provided - """ - must_have = {'eid', 'authors', 'author_ids', 'affiliations', 's2_authors', 's2_author_ids'} - columns = set(df.columns) - if columns & must_have != must_have: - return False, f'The columns {list(must_have - columns)} are missing in `df`' - - return True, None - - - def __verify_slic_df(self, slic_df): - """ - Verify that the given papers dataframe matches the SLIC standard and can be used with Orca - - Parameters - ---------- - slic_df: pandas.DataFrame - A pre-computed DataFrame with SLIC id mappings. - - Returns - ------- - flag: bool - If true, df passes the test and can be used with Orca - error: str, None - If flag is True, None is returned. Otherwise a string with the encountered error is provided - """ - must_have = {'slic_id', 'slic_name', 'scopus_ids', 'scopus_names', 's2_ids', 's2_names'} - columns = set(slic_df.columns) - if columns & must_have != must_have: - return False, f'The columns {list(must_have - columns)} are missing in `slic_df`' - - return True, None - - + + # ───────────────────────────────────────────────────────────────────────── + # Internals + # ───────────────────────────────────────────────────────────────────────── + + def _run_scopus(self, df): + df = df.copy() + for col in ("author_ids", "authors"): + if col not in df.columns: + df[col] = pd.NA + + affiliations_map = self.__generate_affiliations_map(df) + scopus_author_map = self.__generate_author_map(df, "author_ids", "authors") + + duplicates = {} + for entry in self.duplicates: + if not (set(scopus_author_map.keys()) & set(entry)): + continue + + entry = entry.copy() + best_id = sorted(entry, key=lambda x: len(scopus_author_map.get(x, "")), reverse=True)[0] + merged_affiliations = self.__merge_scopus_affiliations(entry, affiliations_map) + affiliations_map[best_id] = merged_affiliations + + entry.remove(best_id) + for x in entry: + if x in scopus_author_map: + del scopus_author_map[x] + if x in affiliations_map: + del affiliations_map[x] + + duplicates[best_id] = list(entry) + + slic_df = { + "slic_id": [], + "slic_name": [], + "scopus_ids": [], + "scopus_names": [], + "scopus_affiliations": [], + "s2_ids": [], + "s2_names": [], + } + + for i, scopus_id in enumerate(scopus_author_map): + scopus_name = scopus_author_map.get(scopus_id) + scopus_affiliations = affiliations_map.get(scopus_id) + if scopus_id in duplicates: + scopus_id = ";".join([scopus_id] + duplicates[scopus_id]) + + slic_df["slic_id"].append(f"S{i}") + slic_df["slic_name"].append(scopus_name) + slic_df["scopus_ids"].append(scopus_id) + slic_df["scopus_names"].append(scopus_name) + slic_df["scopus_affiliations"].append(scopus_affiliations) + slic_df["s2_ids"].append(None) + slic_df["s2_names"].append(None) + + return pd.DataFrame.from_dict(slic_df) + + def _run_s2_only(self, df): + df = df.copy() + for col in ("s2_author_ids", "s2_authors"): + if col not in df.columns: + df[col] = pd.NA + + s2_author_map = self.__generate_author_map(df, "s2_author_ids", "s2_authors") + + visited = set() + groups = [] + + for s2_id in s2_author_map.keys(): + if s2_id in visited: + continue + group = {s2_id} + if s2_id in self.s2_duplicates: + group |= set(self.s2_duplicates[s2_id]) + visited |= group + groups.append(group) + + for root, dups in self.s2_duplicates.items(): + if root not in visited: + group = {root} | set(dups) + visited |= group + groups.append(group) + + rows = { + "slic_id": [], + "slic_name": [], + "scopus_ids": [], + "scopus_names": [], + "scopus_affiliations": [], + "s2_ids": [], + "s2_names": [], + } + + idx = 0 + seen = set() + for g in groups: + name_set = {s2_author_map.get(x, "Unknown") for x in g if x in s2_author_map} + if name_set == {"Unknown"}: + name_set = set() + slic_name = max(name_set, key=len) if name_set else None + + rows["slic_id"].append(f"S{idx}") + rows["slic_name"].append(slic_name) + rows["scopus_ids"].append(None) + rows["scopus_names"].append(None) + rows["scopus_affiliations"].append(None) + rows["s2_ids"].append(";".join(sorted(g))) + rows["s2_names"].append(";".join(sorted(name_set)) if name_set else None) + + seen |= g + idx += 1 + + for s2_id, s2_name in s2_author_map.items(): + if s2_id in seen: + continue + rows["slic_id"].append(f"S{idx}") + rows["slic_name"].append(s2_name if s2_name != "Unknown" else None) + rows["scopus_ids"].append(None) + rows["scopus_names"].append(None) + rows["scopus_affiliations"].append(None) + rows["s2_ids"].append(s2_id) + rows["s2_names"].append(s2_name) + idx += 1 + + slic_df = pd.DataFrame.from_dict(rows) + slic_df = slic_df.loc[slic_df.slic_name.notna()].reset_index(drop=True) + return slic_df + + # ───────────────────────────────────────────────────────────────────────── + # Helpers + # ───────────────────────────────────────────────────────────────────────── + def __add_scopus_duplicates(self, auth_df, duplicates): - """ - Helper function that enriches the results of AuthorMatcher with any previously detected Scopus - duplicates. These duplicates are given a shared S2 id that will be used to connect them in the - succeeding SLIC id creation steps. - - Parameters - ---------- - auth_df: pandas.DataFrame - The results of AuthorMatcher on the working DataFrame - duplicates: list(set), optional - A list of sets where each set contains scopus author ids that refer to the same person. In the ideal case, each - author only has one scopus id. However, this ideal does not hold up in practice and some authors are represented - by two or more scopus ids. Duplicate authors can be found using the Orca.DuplicateAuthorFinder tool. - - Returns - ------- - out_df: pandas.DataFrame - DataFrame that matches the shape of auth_df but contains entries to flag scopus duplicates - """ out_df = pd.DataFrame(columns=auth_df.columns) all_df_ids = set(auth_df.SCOPUS_Author_ID.to_list()) if self.verbose: - print('[Orca]: Scanning for Scopus duplicates in dataset. . .') + print("[Orca]: Scanning for Scopus duplicates in dataset. . .") for scopus_ids in tqdm(duplicates, total=len(duplicates), disable=not self.verbose): if not scopus_ids & all_df_ids: continue @@ -398,534 +451,326 @@ def __add_scopus_duplicates(self, auth_df, duplicates): for scopus_id in tmp_df.SCOPUS_Author_ID.unique(): if row.SCOPUS_Author_ID == scopus_id: continue - else: - new_row = row.copy() - new_row.SCOPUS_Author_ID = scopus_id - out_df = pd.concat([out_df, new_row.to_frame().T], ignore_index=True) + new_row = row.copy() + new_row.SCOPUS_Author_ID = scopus_id + out_df = pd.concat([out_df, new_row.to_frame().T], ignore_index=True) return out_df - - + def __add_s2_duplicates(self, duplicates): - """ - Converts a list of sets into a dictionary such that each key in the dictionary - is an element from a set, and its value is a list of the other elements in that set. - Each set should contain s2 author ids that are known duplicates of each other. No - two pairs of sets can share an author id. All known duplicate of an s2 author id should - be contained within a single set. Sets with a single id (non-duplicates) will be ignored. - - Parameters: - ----------- - duplicates: list(set()) - A list of sets of s2 authors ids to be processed. - - Returns: - -------- - dict: - The processed s2 duplicates which will be resolved in a further processinf step. - - Raises: - ------- - ValueError: - If an id appears in more than one set within the list. - - Example: - -------- - >>> self.__add_s2_duplicates([{1,2}, {3,4}, {5}]) - {1: [2], 2: [1], 3: [4], 4: [3]} - >>> self.__add_s2_duplicates([{1,2}, {2,3}]) - ValueError - """ - out_dict = {} + out = {} seen = set() - for s in duplicates: if any(elem in seen for elem in s): - raise ValueError("Detected multiple entries for s2 duplicates across sets. \ - Make sure that all known duplicates are constrained to a single set") + raise ValueError( + "Detected multiple entries for s2 duplicates across sets. " + "Make sure that all known duplicates are constrained to a single set" + ) seen.update(s) if len(s) == 1: continue - for element in s: - out_dict[element] = [x for x in s if x != element] - return out_dict + out[element] = [x for x in s if x != element] + return out - def __propagate_duplicates(self, a_map, a_duplicates): - """ - Propagate the ids associated with keys in a_map to their duplicate keys. - - For each key in a_map, if duplicate keys exist in a_duplicates, - the associated values from a_map are propagated to these duplicate keys. - The 'a'/'b' notation is used to keep this function generic so that is can be used - to go from s2 --> scopus or scopus --> s2. - - Parameters: - ----------- - a_map: dict - The main author dictionary that is to be updated. - a_duplicates: dict - Dictionary mapping keys to lists of their duplicates. - - Returns: - -------- - None - a_map is modified in place. - """ - a_map_update = {} - - # create an update map based on duplicates for a_id, b_ids in a_map.items(): if a_id in a_duplicates: for dup_a_id in a_duplicates[a_id]: - if dup_a_id not in a_map_update: - a_map_update[dup_a_id] = set() - a_map_update[dup_a_id] |= set(b_ids) - - # convert set to list for each key in the update map + a_map_update.setdefault(dup_a_id, set()).update(set(b_ids)) for dup_a_id in a_map_update: a_map_update[dup_a_id] = list(a_map_update[dup_a_id]) - - # update the main map in place a_map.update(a_map_update) - def __uncouple_author_matches(self, auth_df, s2_duplicates): - """ - Helper function that takes the authors DataFrame produced by AuthorMatcher (could be enriched with scopus duplicates - or not) and finds out which sets of authors ids represent the same individual. This is done by building a graph of - author ids relationships. - - Take for example the following two maps. In this case letters are scopus IDs and numbers are S2 IDs. Each map presents - the relationship between scopus and S2 from the perspective of the key dataset. There are only 2 authors but they are - represented by 2 scopus / 3 S2 IDs for the first author and 1 scopus / 2 S2 IDs for the second author. - - >>> scopus_map = {'A': [1,2], - 'B': [4], - 'C': [3,5]} - - >>> s2_map = {1: ['A'], - 2: ['A'], - 3: ['C'], - 4: ['A'], - 5: ['C']} - - For bigger datasets, these relationships can grow complex and are best modeled by a graph. Both scopus and S2 IDs are - nodes in this graph and their relationship can be modeled with egdes between them. This graph is very disconnected as - there will be many unique authors in any given SLIC dataset. However, weakly connected components of the graph will - signify that all author id nodes in said component belong to the same author. - - Parameters - ---------- - auth_df: pandas.DataFrame - The results of AuthorMatcher on the working DataFrame - s2_duplicates: - - Returns - ------- - matches: list - A list of dictionaries. Each dictionary in the list contains 2 keys: 'scopus' and 's2'. The values are sets of - corresponding scopus/s2 ids - """ - # create the two maps necessary for processing - s2_map = auth_df.groupby('S2_Author_ID')['SCOPUS_Author_ID'].agg(set).to_dict() - scopus_map = auth_df.groupby('SCOPUS_Author_ID')['S2_Author_ID'].agg(set).to_dict() - - # handle s2 duplicates + s2_map = auth_df.groupby("S2_Author_ID")["SCOPUS_Author_ID"].agg(set).to_dict() + scopus_map = auth_df.groupby("SCOPUS_Author_ID")["S2_Author_ID"].agg(set).to_dict() self.__propagate_duplicates(s2_map, s2_duplicates) - - # ensure that no s2 id == scopus id by coincidence - s2_map = {f'B_{k}': {f'A_{x}' for x in v} for k,v in s2_map.items()} - scopus_map = {f'A_{k}': {f'B_{x}' for x in v} for k,v in scopus_map.items()} - # also create a name map which we will use - s2_name_map = auth_df.groupby('S2_Author_ID')['S2_Author_Name'].agg(lambda x: max(x, key=len)).to_dict() + s2_map = {f"B_{k}": {f"A_{x}"} if isinstance(v, str) else {f"A_{x}" for x in v} for k, v in s2_map.items()} + scopus_map = {f"A_{k}": {f"B_{x}"} if isinstance(v, str) else {f"B_{x}" for x in v} for k, v in scopus_map.items()} + + s2_name_map = auth_df.groupby("S2_Author_ID")["S2_Author_Name"].agg(lambda x: max(x, key=len)).to_dict() - # setup the graph G = nx.DiGraph() for k, v_set in scopus_map.items(): for v in v_set: G.add_edge(k, v) - for k, v_set in s2_map.items(): for v in v_set: G.add_edge(k, v) - # get the list of components and process them components = list(nx.weakly_connected_components(G)) matches = [] for component_set in components: - mdict = {'scopus': set(), 's2': set()} + mdict = {"scopus": set(), "s2": set()} for c in component_set: - if c.startswith('A_'): - mdict['scopus'].add(c[2:]) - else: - mdict['s2'].add(c[2:]) - - # get the longest s2 name to use as slic name - str_gen = ((pid, s2_name_map[pid]) for pid in mdict['s2'] if pid in s2_name_map) - _, name = max(str_gen, key=lambda x: len(x[1]), default=(None, 'Unknown')) - mdict['name'] = name + (mdict["scopus"] if c.startswith("A_") else mdict["s2"]).add(c[2:]) + str_gen = ((pid, s2_name_map[pid]) for pid in mdict["s2"] if pid in s2_name_map) + _, name = max(str_gen, key=lambda x: len(x[1]), default=(None, "Unknown")) + mdict["name"] = name matches.append(mdict) return matches - - + def __generate_author_map(self, df, id_col, name_col): - """ - Helper function that generates a map of author ids to author names - - Parameters - ---------- - df: pandas.DataFrame - The SLIC papers DataFrame for which author SLIC ids need to be created - id_col: str - The author ids column. Options are ['author_ids', 's2_author_ids'] - name_col: str - The author names column. Options are ['authors', 's2_authors'] - - Returns - ------- - auth_map: dict - Map where keys are author ids and values are author names - """ if self.verbose: - print(f'[Orca]: Generating {id_col}-{name_col} map. . .') - + print(f"[Orca]: Generating {id_col}-{name_col} map. . .") + if id_col not in df.columns or name_col not in df.columns: + return {} auth_map = {} - tmp_df = df.dropna(subset=[id_col, name_col]) - for id_list, auth_list in tqdm(zip(tmp_df[id_col].to_list(), tmp_df[name_col].to_list()), total=len(tmp_df), disable=not self.verbose): - for auth_id, name in zip(id_list.split(';'), auth_list.split(';')): - if auth_id not in auth_map: + tmp = df[[id_col, name_col]].dropna(how="any") + if tmp.empty: + return {} + for id_list, auth_list in tqdm( + zip(tmp[id_col].to_list(), tmp[name_col].to_list()), total=len(tmp), disable=not self.verbose + ): + if not isinstance(id_list, str) or not isinstance(auth_list, str): + continue + for auth_id, name in zip(id_list.split(";"), auth_list.split(";")): + if auth_id and auth_id not in auth_map: auth_map[auth_id] = name return auth_map - - + def __compute_slic_scopus(self, df): - """ - Helper function applies the SLIC id map to papers that have scopus information - - Parameters - ---------- - df: pandas.DataFrame - The SLIC papers DataFrame for which author SLIC ids need to be created - - Returns - ------- - scopus_df: pandas.DataFrame - papers DataFrame that contains SLIC id and affiliation information - """ - tmp_df = df.loc[~df['eid'].isnull()] # get only scopus papers - - # create maps for scopus author an + tmp_df = df.loc[~df["eid"].isnull()] scopus_authors, scopus_affiliations = {}, {} - for eid, author_ids, affiliations in zip(tmp_df['eid'].to_list(), - tmp_df['author_ids'].to_list(), - tmp_df['affiliations'].to_list()): + + for eid, author_ids, affiliations in zip( + tmp_df["eid"].to_list(), tmp_df["author_ids"].to_list(), tmp_df["affiliations"].to_list() + ): if not pd.isna(author_ids): scopus_authors[eid] = author_ids if not pd.isna(affiliations): if isinstance(affiliations, str): affiliations = ast.literal_eval(affiliations) scopus_affiliations[eid] = affiliations - - scopus_df = { - 'eid': [], - 'slic_author_ids': [], - 'slic_affiliations': [], + + scopus_df = {"eid": [], "slic_author_ids": [], "slic_affiliations": []} + scopus_to_slic = { + x: k + for k, v in zip(self.slic_df.slic_id.to_list(), self.slic_df.scopus_ids.to_list()) + if not pd.isna(v) + for x in v.split(";") } - - # compute map of scopus author id to slic author id - scopus_to_slic = {x: k for k,v in zip(self.slic_df.slic_id.to_list(), self.slic_df.scopus_ids.to_list()) - if not pd.isna(v) for x in v.split(';')} - missing_authors = set() for eid in tmp_df.eid.to_list(): - - slic_author_ids = [] # first replace author_ids information + slic_author_ids = [] author_ids = scopus_authors.get(eid) if author_ids is not None: - for scopus_id in author_ids.split(';'): - scopus_id = str(scopus_id) # should already be string but hard cast to make sure + for scopus_id in author_ids.split(";"): + scopus_id = str(scopus_id) slic_id = scopus_to_slic.get(scopus_id) if slic_id is None: missing_authors.add(scopus_id) else: slic_author_ids.append(str(slic_id)) - aff_dict, del_dict = {}, [] # next update affiliations structure + aff_dict, del_dict = {}, [] affiliations = scopus_affiliations.get(eid) if affiliations is not None: for aff_id, aff_info_shallow in affiliations.items(): if isinstance(aff_info_shallow, list): continue - del_list = [] # items to remove + del_list = [] aff_info = copy.deepcopy(aff_info_shallow) - for i in range(len(aff_info['authors'])): - scopus_id = str(aff_info['authors'][i]) + for i in range(len(aff_info["authors"])): + scopus_id = str(aff_info["authors"][i]) if scopus_id not in scopus_to_slic: del_list.append(scopus_id) missing_authors.add(scopus_id) else: - aff_info['authors'][i] = scopus_to_slic[scopus_id] + aff_info["authors"][i] = scopus_to_slic[scopus_id] for d in del_list: - if d not in aff_info['authors']: - aff_info['authors'].remove(str(d)) - else: - aff_info['authors'].remove(d) + for cand in (d, str(d)): + if cand in aff_info["authors"]: + aff_info["authors"].remove(cand) + break - if not aff_info['authors']: - del_dict.append(aff_id) + if not aff_info["authors"]: + del_dict.append(aff_id) aff_dict[aff_id] = aff_info for d in del_dict: del aff_dict[d] - scopus_df['eid'].append(eid) - if not slic_author_ids: - scopus_df['slic_author_ids'].append(None) - else: - scopus_df['slic_author_ids'].append(";".join(slic_author_ids)) - if not aff_dict: - scopus_df['slic_affiliations'].append(None) - else: - scopus_df['slic_affiliations'].append(aff_dict) + scopus_df["eid"].append(eid) + scopus_df["slic_author_ids"].append(";".join(slic_author_ids) if slic_author_ids else None) + scopus_df["slic_affiliations"].append(aff_dict if aff_dict else None) if len(missing_authors) > 0: - warnings.warn(f'[Orca]: {len(missing_authors)} Scopus IDs did not have corresponding SLIC ID and were removed') - - scopus_df = pd.DataFrame.from_dict(scopus_df) - return scopus_df - - + warnings.warn( + f"[Orca]: {len(missing_authors)} Scopus IDs did not have corresponding SLIC ID and were removed" + ) + + return pd.DataFrame.from_dict(scopus_df) + def __compute_slic_s2(self, df): - """ - Helper function applies the SLIC id map to papers that have S2 information - - Parameters - ---------- - df: pandas.DataFrame - The SLIC papers DataFrame for which author SLIC ids need to be created - - Returns - ------- - s2_df: pandas.DataFrame - papers DataFrame that contains SLIC id and affiliation information - """ - #tmp_df = df.loc[df['eid'].isnull()] # get only s2 papers - tmp_df = df.loc[~df['s2id'].isnull()] - - # compute map of s2 author id to slic author id - s2_to_slic = {x: k for k,v in zip(self.slic_df.slic_id.to_list(), self.slic_df.s2_ids.to_list()) - if not pd.isna(v) for x in v.split(';')} - - s2_df = { - 's2id': [], - 'slic_author_ids': [], + tmp_df = df.loc[~df["s2id"].isnull()] + s2_to_slic = { + x: k + for k, v in zip(self.slic_df.slic_id.to_list(), self.slic_df.s2_ids.to_list()) + if not pd.isna(v) + for x in v.split(";") } - + + s2_df = {"s2id": [], "slic_author_ids": []} missing_authors = set() - for s2id, s2_author_ids in zip(tmp_df['s2id'].to_list(), tmp_df['s2_author_ids'].to_list()): + for s2id, s2_author_ids in zip(tmp_df["s2id"].to_list(), tmp_df["s2_author_ids"].to_list()): slic_author_ids = [] if not pd.isna(s2_author_ids): - for s2_auth_id in s2_author_ids.split(';'): + for s2_auth_id in s2_author_ids.split(";"): if s2_auth_id not in s2_to_slic: missing_authors.add(s2_auth_id) else: slic_author_ids.append(s2_to_slic[s2_auth_id]) - - if slic_author_ids: - s2_df['s2id'].append(s2id) - s2_df['slic_author_ids'].append(";".join(slic_author_ids)) - else: - s2_df['s2id'].append(s2id) - s2_df['slic_author_ids'].append(None) + s2_df["s2id"].append(s2id) + s2_df["slic_author_ids"].append(";".join(slic_author_ids) if slic_author_ids else None) if len(missing_authors) > 0: - warnings.warn(f'[Orca]: {len(missing_authors)} S2 IDs did not have corresponding SLIC ID and were removed') - - s2_df = pd.DataFrame.from_dict(s2_df) - return s2_df + warnings.warn( + f"[Orca]: {len(missing_authors)} S2 IDs did not have corresponding SLIC ID and were removed" + ) + return pd.DataFrame.from_dict(s2_df) def __load_pickle(self, fn): - """ - Helper function for loading pickle files saved in data package - If upgrading to python >=3.9, change this function make use of importlib.resources - - Parameters - ---------- - fn: str - The file name to be loaded - - Returns - ------- - python object stored in the pickle file - """ current_dir = os.path.dirname(os.path.abspath(__file__)) - return pickle.load(open((os.path.join(current_dir, 'data', fn)), 'rb')) - + return pickle.load(open((os.path.join(current_dir, "data", fn)), "rb")) def __generate_affiliations_map(self, df): """ - Helper function for computing a map of affiliations for scopus authors. - The output map is a dict that adheres to the following structure: - { - SCOPUS_AUTHOR_ID: - { - SCOPUS_AFFILIATION_ID: - { - 'name': NAME, # the name of the affiliation - 'country': COUNTRY # country associated with the affiliation - 'first_seen': XXXX # the year when first known paper was published by author with given affiliation - 'last_seen': XXXX # the year when last known paper was published by author with given affiliation - 'papers': XXXX # list of known papers with this affiliation. NOT guaranteed to contain all papers - }, - ... - }, - ... - } - - Parameters - ---------- - df: pandas.DataFrame - The SLIC papers DataFrame for which author SLIC ids need to be created - - Returns - ------- - affiliations_map: dict, - The created map + Build a map SCOPUS_AUTHOR_ID -> affiliation info with paper provenance. + Paper id priority: eid > s2id > doi > synthetic row id. """ + import pandas as pd, ast + + if "affiliations" not in df.columns or "year" not in df.columns: + return {} + + id_priority = ("eid", "s2id", "doi") + pid_col = next((c for c in id_priority if c in df.columns), None) + if pid_col is None: + df = df.copy() + df["_orca_row_id"] = [f"row_{i}" for i in range(len(df))] + pid_col = "_orca_row_id" + affiliations_map = {} - for eid, year, affiliations in zip(df.eid.to_list(), df.year.to_list(), df.affiliations.to_list()): - - # handle missing / unconverted affiliations + pid_series = df[pid_col].astype(object).where(pd.notna(df[pid_col]), None) + year_series = df["year"] + aff_series = df["affiliations"] + + for pid, year, affiliations in zip(pid_series.tolist(), year_series.tolist(), aff_series.tolist()): if pd.isna(affiliations): continue if isinstance(affiliations, str): affiliations = ast.literal_eval(affiliations) - - # get the year - if pd.isna(year): - year = 0 - else: - year = int(year) - + + year = 0 if pd.isna(year) else int(year) + paper_id = None if pid is None else str(pid) + for aff_id, info in affiliations.items(): if isinstance(info, list): continue - aff_name = info.get('name', 'Unknown') - aff_country = info.get('country', 'Unknown') + aff_name = info.get("name", "Unknown") + aff_country = info.get("country", "Unknown") - for auth_id in info.get('authors', []): + for auth_id in info.get("authors", []): if auth_id not in affiliations_map: affiliations_map[auth_id] = {} if aff_id not in affiliations_map[auth_id]: - affiliations_map[auth_id][aff_id] = {} - first_seen = affiliations_map[auth_id][aff_id].get('first_seen', 1*10**4) - last_seen = affiliations_map[auth_id][aff_id].get('last_seen', -1) - - if 'name' not in affiliations_map[auth_id][aff_id]: - affiliations_map[auth_id][aff_id]['name'] = aff_name - affiliations_map[auth_id][aff_id]['country'] = aff_country - affiliations_map[auth_id][aff_id]['first_seen'] = year - affiliations_map[auth_id][aff_id]['last_seen'] = year - affiliations_map[auth_id][aff_id]['papers'] = {eid} - else: - affiliations_map[auth_id][aff_id]['papers'].add(eid) - if year < first_seen or first_seen == 0: - affiliations_map[auth_id][aff_id]['first_seen'] = year - elif year > last_seen or last_seen == 0: - affiliations_map[auth_id][aff_id]['last_seen'] = year + affiliations_map[auth_id][aff_id] = { + "name": aff_name, + "country": aff_country, + "first_seen": year, + "last_seen": year, + "papers": set(), + } + + entry = affiliations_map[auth_id][aff_id] + if paper_id is not None: + entry["papers"].add(paper_id) + if entry["first_seen"] in (0, None) or (year and year < entry["first_seen"]): + entry["first_seen"] = year + if entry["last_seen"] in (0, None) or (year and year > entry["last_seen"]): + entry["last_seen"] = year - # handle the missing years for auth_id in affiliations_map: for aff_id, aff_info in affiliations_map[auth_id].items(): - if not aff_info['first_seen']: - affiliations_map[auth_id][aff_id]['first_seen'] = None - if not aff_info['last_seen']: - affiliations_map[auth_id][aff_id]['last_seen'] = None - return affiliations_map + aff_info["papers"] = list(aff_info["papers"]) + if not aff_info["first_seen"]: + aff_info["first_seen"] = None + if not aff_info["last_seen"]: + aff_info["last_seen"] = None + return affiliations_map def __merge_scopus_affiliations(self, scopus_ids, data): - """ - Helper function for merging the scopus affiliation maps for entries where multiple scopus - ids correspond to the same author - - Parameters - ---------- - scopus_ids: list - A list of scopus author ids - data: dict - The scopus affiliations map - - Returns - ------- - merged: dict - A dictionary that is a result of merging the dictionaries associated with the - duplicate author ids. For shared items, 'first_seen' is the earliest and 'last_seen' is - the latest among all the dictionaries. For unshared items, their details are kept - as is. If 'first_seen' or 'last_seen' is 'Unknown' in one dictionary but has a valid - integer value in another, the integer value is considered. If no valid integer value - exists for 'first_seen' or 'last_seen', 'Unknown' is set as the value. - """ merged = {} - all_keys = set(k for key in scopus_ids if key in data for k in data[key].keys()) - + if not data or not scopus_ids: + return None + + all_keys = set() + for sid in scopus_ids: + if sid in data: + all_keys.update(data[sid].keys()) + for key in all_keys: - items = [data[k][key] for k in scopus_ids if k in data and key in data[k]] - names = [item['name'] for item in items if item['name'] != 'Unknown'] - countries = [item['country'] for item in items if item['country'] != 'Unknown'] - papers = list(set.union(*[item['papers'] for item in items])) + items = [data[sid][key] for sid in scopus_ids if sid in data and key in data[sid]] + if not items: + continue + + names = [it.get("name", "Unknown") for it in items if it.get("name") not in (None, "Unknown")] + countries = [it.get("country", "Unknown") for it in items if it.get("country") not in (None, "Unknown")] + + paper_sets = [] + for it in items: + p = it.get("papers", []) + if isinstance(p, set): + paper_sets.append(p) + elif p is None: + continue + else: + paper_sets.append(set(p)) + merged_papers = sorted(set().union(*paper_sets)) if paper_sets else [] + + first_vals = [it.get("first_seen") for it in items if isinstance(it.get("first_seen"), int)] + last_vals = [it.get("last_seen") for it in items if isinstance(it.get("last_seen"), int)] + merged[key] = { - 'name': names[0] if names else 'Unknown', # use first valid affiliation name, or 'Unknown' if no valid names - 'country': countries[0] if countries else 'Unknown', # use first valid affiliation country, or 'Unknown' if no valid countries - 'first_seen': min((item['first_seen'] for item in items if isinstance(item['first_seen'], int)), default='Unknown'), - 'last_seen': max((item['last_seen'] for item in items if isinstance(item['last_seen'], int)), default='Unknown'), - 'papers': papers + "name": names[0] if names else "Unknown", + "country": countries[0] if countries else "Unknown", + "first_seen": min(first_vals) if first_vals else "Unknown", + "last_seen": max(last_vals) if last_vals else "Unknown", + "papers": merged_papers, } - merged = merged if merged else None # {} --> None - return merged - - - ### GETTERS / SETTERS - - + + return merged or None + + # ───────────────────────────────────────────────────────────────────────── + # Getters / Setters + # ───────────────────────────────────────────────────────────────────────── + @property def duplicates(self): return self._duplicates - - @property - def s2_duplicates(self): - return self._s2_duplicates @duplicates.setter def duplicates(self, duplicates): if duplicates is None: self._duplicates = [] elif isinstance(duplicates, list): - self._duplicates = duplicates - else: - raise TypeError(f' {type(duplicates)} is an invalid type for `duplicates`') - - """ - Code where we have a precomputed duplicates. - @duplicates.setter - def duplicates(self, duplicates): - if duplicates is None: - self._duplicates = self.__load_pickle(self.DUPLICATES_1M) - elif isinstance(duplicates, list): - self._duplicates = self.__load_pickle(self.DUPLICATES_1M) + duplicates + self._duplicates = duplicates else: - raise TypeError(f' {type(duplicates)} is an invalid type for `duplicates`') - """ - - + raise TypeError(f"{type(duplicates)} is an invalid type for `duplicates`") + + @property + def s2_duplicates(self): + return self._s2_duplicates + @s2_duplicates.setter def s2_duplicates(self, s2_duplicates): if s2_duplicates is None: @@ -933,4 +778,4 @@ def s2_duplicates(self, s2_duplicates): elif isinstance(s2_duplicates, list): self._s2_duplicates = self.__add_s2_duplicates(s2_duplicates) else: - raise TypeError(f' {type(s2_duplicates)} is an invalid type for `s2_duplicates`') \ No newline at end of file + raise TypeError(f"{type(s2_duplicates)} is an invalid type for `s2_duplicates`") diff --git a/TELF/pre_processing/Squirrel/__init__.py b/TELF/pre_processing/Squirrel/__init__.py old mode 100755 new mode 100644 diff --git a/TELF/pre_processing/Vulture/modules/simple_clean.py b/TELF/pre_processing/Vulture/modules/simple_clean.py old mode 100755 new mode 100644 diff --git a/TELF/pre_processing/Vulture/vulture.py b/TELF/pre_processing/Vulture/vulture.py old mode 100755 new mode 100644 diff --git a/TELF/pre_processing/iPenguin/scripts/sync.sh b/TELF/pre_processing/iPenguin/scripts/sync.sh old mode 100755 new mode 100644 diff --git a/TELF/version.py b/TELF/version.py index 59120223..fa571767 100644 --- a/TELF/version.py +++ b/TELF/version.py @@ -1 +1 @@ -__version__ = "0.0.43" +__version__ = "0.0.44" diff --git a/data/sample_terms3.md b/data/sample_terms3.md old mode 100755 new mode 100644 diff --git a/docs/.buildinfo b/docs/.buildinfo new file mode 100644 index 00000000..29b0e15a --- /dev/null +++ b/docs/.buildinfo @@ -0,0 +1,4 @@ +# Sphinx build info version 1 +# This file records the configuration used when building these files. When it is not found, a full rebuild will be done. +config: 8541cce001b52ed2e074c33e32340b3e +tags: 645f666f9bcd5a90fca523b33c5a78b7 diff --git a/docs/ArcticFox.html b/docs/ArcticFox.html index fafa1b87..0ecbdc87 100644 --- a/docs/ArcticFox.html +++ b/docs/ArcticFox.html @@ -8,7 +8,7 @@ - TELF.post_processing.ArcticFox: Report generation tool for text data from HNMFk using local LLMs — TELF 0.0.43 documentation + TELF.post_processing.ArcticFox: Report generation tool for text data from HNMFk using local LLMs — TELF 0.0.44 documentation @@ -16,31 +16,28 @@ document.documentElement.dataset.mode = localStorage.getItem("mode") || ""; document.documentElement.dataset.theme = localStorage.getItem("theme") || ""; - - - - + + + + + + + + + - + - - - - + + + - + @@ -51,7 +48,6 @@ - @@ -67,8 +63,19 @@ Back to top - - + + + + + + +
+
+
@@ -76,6 +83,7 @@ Ctrl+K -
-
+ +
@@ -100,8 +108,7 @@ - -
+
@@ -126,16 +133,20 @@ -

TELF 0.0.43 documentation

+

TELF 0.0.44 documentation

-
+