cnhkmcp 2.1.9__py3-none-any.whl → 2.2.0__py3-none-any.whl

cnhkmcp/untracked/skills/brain-feature-implementation/scripts/fetch_dataset.py

@@ -0,0 +1,107 @@
+import json
+import os
+import argparse
+import pandas as pd
+from pathlib import Path
+import sys
+
+print("Script started...", flush=True)
+
+# Ensure local imports work by adding the script directory to sys.path
+script_dir = Path(__file__).resolve().parent
+sys.path.append(str(script_dir))
+
+try:
+    import ace_lib
+except ImportError:
+    print("Error: Could not import 'ace_lib'. Make sure it is in the same directory.")
+    sys.exit(1)
+
+def load_config(config_path):
+    try:
+        with open(config_path, 'r') as f:
+            return json.load(f)
+    except Exception as e:
+        print(f"Error loading config file: {e}")
+        return None
+
+def main():
+    parser = argparse.ArgumentParser(description="Fetch dataset fields from WorldQuant BRAIN")
+    parser.add_argument("--datasetid", required=True, help="ID of the dataset to fetch (e.g., specific dataset ID)")
+    parser.add_argument("--region", default="USA", help="Region (default: USA)")
+    parser.add_argument("--delay", type=int, default=1, help="Delay (default: 1)")
+    parser.add_argument("--universe", default="TOP3000", help="Universe (default: TOP3000)")
+    parser.add_argument("--instrument-type", default="EQUITY", dest="instrument_type", help="Instrument Type (default: EQUITY)")
+
+    args = parser.parse_args()
+
+    # Determine paths relative to this script
+    # User requested: robust and no absolute paths hardcoded
+    workspace_dir = script_dir.parent
+    config_path = workspace_dir / "config.json"
+    data_dir = workspace_dir / "data"
+
+    # Ensure data directory exists
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load configuration
+    if not config_path.exists():
+        print(f"Error: Config file not found at {config_path}")
+        sys.exit(1)
+
+    config = load_config(config_path)
+    if not config:
+        sys.exit(1)
+
+    # Extract credentials
+    creds = config.get("BRAIN_CREDENTIALS", {})
+    email = creds.get("email")
+    password = creds.get("password")
+
+    if not email or not password:
+        print("Error: BRAIN_CREDENTIALS (email/password) not found in config.json")
+        sys.exit(1)
+
+    # Override ace_lib.get_credentials to use our config values
+    # ace_lib.start_session() internally calls get_credentials()
+    ace_lib.get_credentials = lambda: (email, password)
+
+    try:
+        print(f"Logging in as {email}...")
+        session = ace_lib.start_session()
+        
+        print(f"Fetching datafields for dataset: {args.datasetid} (Region: {args.region}, Delay: {args.delay})...")
+        
+        # Fetch datafields using the library function
+        df = ace_lib.get_datafields(
+            session, 
+            dataset_id=args.datasetid, 
+            region=args.region, 
+            delay=args.delay,
+            universe=args.universe,
+            instrument_type=args.instrument_type
+        )
+
+        if df is None or df.empty:
+            print("Warning: No data found or empty response.")
+        else:
+            # Construct a safe filename and folder name
+            safe_dataset_id = "".join([c for c in args.datasetid if c.isalnum() or c in ('-','_')])
+            folder_name = f"{safe_dataset_id}_{args.region}_delay{args.delay}"
+            dataset_folder = data_dir / folder_name
+            dataset_folder.mkdir(parents=True, exist_ok=True)
+            
+            filename = f"{folder_name}.csv"
+            output_path = dataset_folder / filename
+            
+            print(f"Saving {len(df)} records to {output_path}...")
+            df.to_csv(output_path, index=False)
+            print("Success.")
+
+    except Exception as e:
+        print(f"An error occurred during execution: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    main()

cnhkmcp/untracked/skills/brain-feature-implementation/scripts/helpful_functions.py

@@ -0,0 +1,180 @@
+import json
+import os
+from typing import Union
+
+import pandas as pd
+from pandas.io.formats.style import Styler
+
+brain_api_url = os.environ.get("BRAIN_API_URL", "https://api.worldquantbrain.com")
+brain_url = os.environ.get("BRAIN_URL", "https://platform.worldquantbrain.com")
+
+
+def make_clickable_alpha_id(alpha_id: str) -> str:
+    """
+    Create a clickable HTML link for an alpha ID.
+
+    Args:
+        alpha_id (str): The ID of the alpha.
+
+    Returns:
+        str: An HTML string containing a clickable link to the alpha's page on the platform.
+    """
+
+    url = brain_url + "/alpha/"
+    return f'<a href="{url}{alpha_id}">{alpha_id}</a>'
+
+
+def prettify_result(
+    result: list, detailed_tests_view: bool = False, clickable_alpha_id: bool = False
+) -> Union[pd.DataFrame, Styler]:
+    """
+    Combine and format simulation results into a single DataFrame for analysis.
+
+    Args:
+        result (list): A list of dictionaries containing simulation results.
+        detailed_tests_view (bool, optional): If True, include detailed test results. Defaults to False.
+        clickable_alpha_id (bool, optional): If True, make alpha IDs clickable. Defaults to False.
+
+    Returns:
+        pandas.DataFrame or pandas.io.formats.style.Styler: A DataFrame containing formatted results,
+        optionally with clickable alpha IDs.
+    """
+    list_of_is_stats = [result[x]["is_stats"] for x in range(len(result)) if result[x]["is_stats"] is not None]
+    is_stats_df = pd.concat(list_of_is_stats).reset_index(drop=True)
+    is_stats_df = is_stats_df.sort_values("fitness", ascending=False)
+
+    expressions = {
+        result[x]["alpha_id"]: (
+            {
+                "selection": result[x]["simulate_data"]["selection"],
+                "combo": result[x]["simulate_data"]["combo"],
+            }
+            if result[x]["simulate_data"]["type"] == "SUPER"
+            else result[x]["simulate_data"]["regular"]
+        )
+        for x in range(len(result))
+        if result[x]["is_stats"] is not None
+    }
+    expression_df = pd.DataFrame(list(expressions.items()), columns=["alpha_id", "expression"])
+
+    list_of_is_tests = [result[x]["is_tests"] for x in range(len(result)) if result[x]["is_tests"] is not None]
+    is_tests_df = pd.concat(list_of_is_tests, sort=True).reset_index(drop=True)
+    is_tests_df = is_tests_df[is_tests_df["result"] != "WARNING"]
+    if detailed_tests_view:
+        cols = ["limit", "result", "value"]
+        is_tests_df["details"] = is_tests_df[cols].to_dict(orient="records")
+        is_tests_df = is_tests_df.pivot(index="alpha_id", columns="name", values="details").reset_index()
+    else:
+        is_tests_df = is_tests_df.pivot(index="alpha_id", columns="name", values="result").reset_index()
+
+    alpha_stats = pd.merge(is_stats_df, expression_df, on="alpha_id")
+    alpha_stats = pd.merge(alpha_stats, is_tests_df, on="alpha_id")
+    alpha_stats = alpha_stats.drop(columns=alpha_stats.columns[(alpha_stats == "PENDING").any()])
+    alpha_stats.columns = alpha_stats.columns.str.replace("(?<=[a-z])(?=[A-Z])", "_", regex=True).str.lower()
+    if clickable_alpha_id:
+        return alpha_stats.style.format({"alpha_id": lambda x: make_clickable_alpha_id(str(x))})
+    return alpha_stats
+
+
+def concat_pnl(result: list) -> pd.DataFrame:
+    """
+    Combine PnL results from multiple alphas into a single DataFrame.
+
+    Args:
+        result (list): A list of dictionaries containing simulation results with PnL data.
+
+    Returns:
+        pandas.DataFrame: A DataFrame containing combined PnL data for all alphas.
+    """
+    list_of_pnls = [result[x]["pnl"] for x in range(len(result)) if result[x]["pnl"] is not None]
+    pnls_df = pd.concat(list_of_pnls).reset_index()
+
+    return pnls_df
+
+
+def concat_is_tests(result: list) -> pd.DataFrame:
+    """
+    Combine in-sample test results from multiple alphas into a single DataFrame.
+
+    Args:
+        result (list): A list of dictionaries containing simulation results with in-sample test data.
+
+    Returns:
+        pandas.DataFrame: A DataFrame containing combined in-sample test results for all alphas.
+    """
+    is_tests_list = [result[x]["is_tests"] for x in range(len(result)) if result[x]["is_tests"] is not None]
+    is_tests_df = pd.concat(is_tests_list, sort=True).reset_index(drop=True)
+    return is_tests_df
+
+
+def save_simulation_result(result: dict) -> None:
+    """
+    Save the simulation result to a JSON file in the 'simulation_results' folder.
+
+    Args:
+        result (dict): A dictionary containing the simulation result for an alpha.
+    """
+
+    alpha_id = result["id"]
+    region = result["settings"]["region"]
+    folder_path = "simulation_results/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}")
+
+    os.makedirs(folder_path, exist_ok=True)
+
+    with open(file_path, "w", encoding="utf-8") as file:
+        json.dump(result, file)
+
+
+def save_pnl(pnl_df: pd.DataFrame, alpha_id: str, region: str) -> None:
+    """
+    Save the PnL data for an alpha to a CSV file in the 'alphas_pnl' folder.
+
+    Args:
+        pnl_df (pandas.DataFrame): The DataFrame containing PnL data.
+        alpha_id (str): The ID of the alpha.
+        region (str): The region for which the PnL data was generated.
+    """
+
+    folder_path = "alphas_pnl/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv")
+    os.makedirs(folder_path, exist_ok=True)
+
+    pnl_df.to_csv(file_path)
+
+
+def save_yearly_stats(yearly_stats: pd.DataFrame, alpha_id: str, region: str):
+    """
+    Save the yearly statistics for an alpha to a CSV file in the 'yearly_stats' folder.
+
+    Args:
+        yearly_stats (pandas.DataFrame): The DataFrame containing yearly statistics.
+        alpha_id (str): The ID of the alpha.
+        region (str): The region for which the statistics were generated.
+    """
+
+    folder_path = "yearly_stats/"
+    file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv")
+    os.makedirs(folder_path, exist_ok=True)
+
+    yearly_stats.to_csv(file_path, index=False)
+
+
+def expand_dict_columns(data: pd.DataFrame) -> pd.DataFrame:
+    """
+    Expand dictionary columns in a DataFrame into separate columns.
+
+    Args:
+        data (pandas.DataFrame): The input DataFrame with dictionary columns.
+
+    Returns:
+        pandas.DataFrame: A new DataFrame with expanded columns.
+    """
+    dict_columns = list(filter(lambda x: isinstance(data[x].iloc[0], dict), data.columns))
+    new_columns = pd.concat(
+        [data[col].apply(pd.Series).rename(columns=lambda x: f"{col}_{x}") for col in dict_columns],
+        axis=1,
+    )
+
+    data = pd.concat([data, new_columns], axis=1)
+    return data

cnhkmcp/untracked/skills/brain-feature-implementation/scripts/implement_idea.py

@@ -0,0 +1,164 @@
+import pandas as pd
+from pathlib import Path
+import argparse
+import sys
+import re
+import json
+import time
+
+def load_data(dataset_name=None):
+    script_dir = Path(__file__).resolve().parent
+    workspace_dir = script_dir.parent
+    
+    if not dataset_name:
+        data_root = workspace_dir / "data"
+        if not data_root.exists():
+            print("Error: Data directory not found.", file=sys.stderr)
+            sys.exit(1)
+            
+        subdirs = [d for d in data_root.iterdir() if d.is_dir()]
+        
+        if len(subdirs) == 1:
+            dataset_name = subdirs[0].name
+            print(f"Auto-detected dataset: {dataset_name}", file=sys.stderr)
+        elif len(subdirs) > 1:
+            print("Error: Multiple datasets found. Please specify --dataset.", file=sys.stderr)
+            print("Available datasets:", file=sys.stderr)
+            for d in subdirs:
+                print(f"  {d.name}", file=sys.stderr)
+            sys.exit(1)
+        else:
+            print("Error: No dataset folders found inside data directory.", file=sys.stderr)
+            sys.exit(1)
+
+    dataset_dir = workspace_dir / "data" / dataset_name
+    data_path = dataset_dir / f"{dataset_name}.csv"
+    
+    print(f"Loading data from {data_path}...", file=sys.stderr)
+    try:
+        df = pd.read_csv(data_path)
+        return df, dataset_dir
+    except FileNotFoundError:
+        print(f"Error: Data file not found at {data_path}. Please run fetch_dataset.py first.", file=sys.stderr)
+        sys.exit(1)
+
+def extract_keys_from_template(template):
+    return re.findall(r'\{([A-Za-z0-9_]+)\}', template)
+
+def match_single_horizon_auto(df, template):
+    """
+    Auto-detects metrics from template and finds matching fields.
+    """
+    metrics = extract_keys_from_template(template)
+    if not metrics:
+        print("Error: No variables found in template (use {variable} format).", file=sys.stderr)
+        return []
+
+    # Sort metrics by length descending to match most specific suffixes first
+    metrics = sorted(metrics, key=len, reverse=True)
+    primary = metrics[0]
+    
+    # Try different separators or exact match
+    # we look for columns that end with the primary metric, optionally followed by numeric suffix (e.g. _1234)
+    # Regex: .*<primary>(?:_\d+)?$
+    import re
+    primary_regex = re.escape(primary) + r'(?:_\d+)?$'
+    candidates = df[df['id'].str.match(f'.*{primary_regex}')]['id'].unique().tolist()
+    
+    results = []
+    seen = set()
+    
+    # Try different separators or exact match
+    # We look for columns that contain the primary metric at any position
+    import re
+    primary_regex = re.escape(primary)
+    candidates = df[df['id'].str.contains(primary_regex, regex=True)]['id'].unique().tolist()
+    
+    results = []
+    seen = set()
+    
+    for cand in candidates:
+        # Determine base prefix
+        # We identify the prefix by taking everything before the first occurrence of the primary metric
+        match = re.search(re.escape(primary), cand)
+        if not match:
+             continue
+        
+        # Base includes everything up to the metric (e.g., "dataset_prefix_")
+        base = cand[:match.start()]
+
+        # Verify other metrics exist with this base
+        field_map = {primary: cand}
+        all_found = True
+        
+        for m in metrics[1:]:
+            # Construct target pattern for other metrics: Must start with the same base followed by the metric
+            # We allow any suffix after the metric (e.g. IDs, versions)
+            target_pattern = f"^{re.escape(base)}{re.escape(m)}"
+            target_matches = df[df['id'].str.match(target_pattern)]['id'].tolist()
+            
+            if not target_matches:
+                all_found = False
+                break
+            # Use the first match found for the secondary metric
+            field_map[m] = target_matches[0]
+            
+        if all_found:
+            try:
+                expr = template.format(**field_map)
+                if expr not in seen:
+                    seen.add(expr)
+                    # Create a readable label for the horizon/group
+                    if base:
+                        # Strip standard separators
+                        horizon_label = base.strip("_")
+                    else:
+                        horizon_label = "global"
+                        
+                    results.append((horizon_label, expr))
+            except KeyError as e:
+                continue
+                
+    return results
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate Alpha Expressions based on patterns")
+    parser.add_argument("--template", required=True, help="Python format string (e.g. '{st_dev} / abs({mean})')")
+    parser.add_argument("--dataset", help="Name of the dataset folder. Auto-detected if only one exists.")
+    
+    args = parser.parse_args()
+    
+    df, dataset_dir = load_data(args.dataset)
+    
+    results = match_single_horizon_auto(df, args.template)
+        
+    # Output
+    expression_list = []
+    if not results:
+        print("No matching expressions found.")
+    else:
+        print(f"Generated {len(results)} expressions:\n")
+        # print(f"{'Context':<30} | Expression")
+        # print("-" * 120)
+        
+        for context, expr in results:
+            # print(f"{context:<30} | {expr}")
+            expression_list.append(expr)
+            
+    # Save results to JSON (Always save for debugging)
+    timestamp = int(time.time())
+    json_output = {
+        "template": args.template,
+        "expression_list": expression_list
+    }
+    
+    output_file = dataset_dir / f"idea_{timestamp}.json"
+    try:
+        with open(output_file, 'w') as f:
+            json.dump(json_output, f, indent=4)
+        print(f"\nSaved idea configuration to: {output_file}")
+    except Exception as e:
+        print(f"Error saving JSON: {e}", file=sys.stderr)
+
+if __name__ == "__main__":
+    main()

cnhkmcp/untracked/skills/brain-feature-implementation/scripts/merge_expression_list.py

@@ -0,0 +1,88 @@
+import json
+import argparse
+from pathlib import Path
+import sys
+
+def load_data_dir(dataset_name=None):
+    script_dir = Path(__file__).resolve().parent
+    workspace_dir = script_dir.parent
+    
+    if not dataset_name:
+        data_root = workspace_dir / "data"
+        if not data_root.exists():
+            print("Error: Data directory not found.", file=sys.stderr)
+            sys.exit(1)
+            
+        subdirs = [d for d in data_root.iterdir() if d.is_dir()]
+        
+        if len(subdirs) == 1:
+            dataset_name = subdirs[0].name
+            print(f"Auto-detected dataset: {dataset_name}", file=sys.stderr)
+            return workspace_dir / "data" / dataset_name
+        elif len(subdirs) > 1:
+            print("Error: Multiple datasets found. Please specify --dataset.", file=sys.stderr)
+            sys.exit(1)
+        else:
+            print("Error: No dataset folders found inside data directory.", file=sys.stderr)
+            sys.exit(1)
+
+    return workspace_dir / "data" / dataset_name
+
+def main():
+    parser = argparse.ArgumentParser(description="Merge all generated expressions from idea JSON files.")
+    parser.add_argument("--dataset", help="Name of the dataset folder containing idea JSONs.")
+    parser.add_argument("--output", default="final_expressions.json", help="Output filename.")
+    
+    args = parser.parse_args()
+    
+    dataset_dir = load_data_dir(args.dataset)
+    
+    if not dataset_dir.exists():
+        print(f"Error: Dataset directory {dataset_dir} does not exist.", file=sys.stderr)
+        sys.exit(1)
+        
+    all_expressions = []
+    
+    # Find all idea_*.json files
+    json_files = list(dataset_dir.glob("idea_*.json"))
+    
+    if not json_files:
+        print(f"No idea_*.json files found in {dataset_dir}", file=sys.stderr)
+        sys.exit(0)
+        
+    print(f"Found {len(json_files)} idea files. Merging...")
+    
+    for jf in json_files:
+        try:
+            with open(jf, 'r') as f:
+                data = json.load(f)
+                exprs = data.get("expression_list", [])
+                if exprs:
+                    all_expressions.extend(exprs)
+                    print(f"  + {jf.name}: {len(exprs)} expressions")
+                else:
+                    print(f"  - {jf.name}: 0 expressions")
+        except Exception as e:
+            print(f"  ! Error reading {jf.name}: {e}", file=sys.stderr)
+            
+    # Remove duplicates if desired? Usually we keep them or set them. 
+    # Let's make unique to be safe, but preserve order as best as possible.
+    unique_expressions = []
+    seen = set()
+    for ex in all_expressions:
+        if ex not in seen:
+            unique_expressions.append(ex)
+            seen.add(ex)
+            
+    output_path = dataset_dir / args.output
+    
+    try:
+        with open(output_path, 'w') as f:
+            json.dump(unique_expressions, f, indent=4)
+        print(f"\nSuccessfully merged {len(unique_expressions)} unique expressions.")
+        print(f"Output saved to: {output_path}")
+    except Exception as e:
+        print(f"Error saving output: {e}", file=sys.stderr)
+
+if __name__ == "__main__":
+    main()