eic_for_heuristic.py

import re
import os
import json
import yaml
import time
import signal
import socket
import random
import logging
import datetime
import traceback
import numpy as np
import pandas as pd
from pathlib import Path
from socket import gethostname
from argparse import ArgumentParser
from setproctitle import setproctitle
import src.nd2py.nd2py as nd
from src.nd2py.nd2py.utils import seed_all, init_logger, RMSE_score, R2_score

_logger = logging.getLogger('src')


def search(args):
    if '=' in args.function:
        f = nd.parse(args.function.split('=', 1)[1])
        binary = list(set(op.__class__ for op in f.iter_preorder() if op.n_operands == 2))
        unary = list(set(op.__class__ for op in f.iter_preorder() if op.n_operands == 1))
        leaf = list(set(op for op in f.iter_preorder() if isinstance(op, nd.Number)))
        variables = list(set(op.name for op in f.iter_preorder() if isinstance(op, nd.Variable)))
        X = {var: np.random.uniform(-5, 5, (args.sample_num,)) for var in variables}
        y = f.eval(X)
        log = {
            'target function': args.function,
            'binary operators': [op.__name__ for op in binary],
            'unary operators': [op.__name__ for op in unary],
            'leaf': [op.to_str(number_format=".2f") for op in leaf],
            'variables': list(X.keys()),
        }
    else:
        from data.pmlb import pmlb
        _logger.info(f'fetching {args.function} from PMLB...')
        df = pmlb.fetch_data(args.function, local_cache_dir='./data/pmlb/datasets')
        if df.shape[0] > args.sample_num: 
            df = df.sample(args.sample_num, random_state=args.seed)
        else: 
            args.sample_num = df.shape[0]
        _logger.info(f'Done, df.shape = {df.shape}')
        X = {col:df[col].values for col in df.columns}
        y = X.pop('target')
        binary = [nd.Mul, nd.Div, nd.Add, nd.Sub]
        unary = [nd.Sqrt, nd.Cos, nd.Sin, nd.Pow2, nd.Pow3, nd.Exp, nd.Inv, nd.Neg, nd.Arcsin, nd.Arccos, nd.Cot, nd.Log, nd.Tanh]
        leaf = [nd.Number(1), nd.Number(2), nd.Number(np.pi)]
        log = {
            'target function': args.function,
            'binary operators': [op.__name__ for op in binary],
            'unary operators': [op.__name__ for op in unary],
            'leaf': [op.to_str(number_format=".2f") for op in leaf],
            'variables': list(X.keys()),
        }
        try:
            with open(f'./data/pmlb/datasets/{args.function}/metadata.yaml', 'r') as f:
                metadata = yaml.load(f, Loader=yaml.Loader)['description']
            metadata = [l.strip() for l in metadata.split('\n')]
            target, eq = metadata[metadata.index('')+1].split(' = ', 1)
            eq = nd.parse(eq)
            log['target function'] = log['target function'] + ' ({} = {})'.format(target, eq.to_str(number_format=".2f"))
            if args.cheat:
                binary = list(set(op.__class__ for op in eq.iter_preorder() if op.n_operands == 2))
                unary = list(set(op.__class__ for op in eq.iter_preorder() if op.n_operands == 1))
                leaf = list(set(op for op in eq.iter_preorder() if isinstance(op, nd.Number)))
                log['binary operators'] = [op.__name__ for op in binary]
                log['unary operators'] = [op.__name__ for op in unary]
                log['leaf'] = [op.to_str(number_format=".2f") for op in leaf]
        except Exception as e:
            _logger.warning(e)
    _logger.note('\n'.join(f'{k}: {v if not isinstance(v, list) else "[" + ", ".join(v) + "]"}' for k, v in log.items()))

    y_raw = y.copy()
    if args.target_noise > 0:
        y = y + args.target_noise * np.random.normal(0, y.std(), size=y.shape)
        _logger.note(f'Added Gaussian noise to target with std = {args.target_noise}')

    if args.method == 'eic-mcts':
        from src.heuristic.mcts.mcts import MCTS
        est = MCTS(
            n_iter=args.n_iter,
            keep_vars=args.keep_vars,
            normalize_y=args.normalize_y,
            normalize_X=args.normalize_X,
            remove_abnormal=args.remove_abnormal,
            binary=binary,
            unary=unary,
            leaf=leaf,
            log_per_sec=args.log_per_sec,
            max_len=args.max_len,
            save_path=os.path.join(args.save_dir, 'records.json'),
            max_var=10,
            random_state=args.seed,
            time_limit=args.time_limit,
            eta=args.eta,
            c=args.c,
            alpha=args.alpha,
            ratio=args.ratio,
        )
    elif args.method == 'eic-gp':
        from src.heuristic.gp.gp import MCTS
        est = MCTS(
            n_iter=args.n_iter,
            keep_vars=args.keep_vars,
            normalize_y=args.normalize_y,
            normalize_X=args.normalize_X,
            remove_abnormal=args.remove_abnormal,
            binary=binary,
            unary=unary,
            leaf=leaf,
            log_per_sec=args.log_per_sec,
            max_len=args.max_len,
            save_path=os.path.join(args.save_dir, 'records.json'),
            max_var=10,
            random_state=args.seed,
            time_limit=args.time_limit,
            eta=args.eta,
            c=args.c,
            alpha=args.alpha,
            ratio=args.ratio,
        )
    else:
        raise ValueError(f'Unknown method: {args.method}')

    try:
        def early_stop(r2, complexity, eq):
            if args.target_noise:
                r2 = 1 - np.mean((y_raw - eq.eval(X)) ** 2) / np.var(y_raw)
            return r2 > 0.99999
        status = est.fit(X, y, use_tqdm=False, early_stop=early_stop)
        _logger.info('Finished')
    except KeyboardInterrupt:
        status = 'interrupted'
        _logger.note('Interrupted')
    except Exception:
        status = 'error'
        _logger.error(traceback.format_exc())

    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        y_pred = est.predict(X)
        rmse = RMSE_score(y_raw, y_pred)
        r2 = R2_score(y_raw, y_pred)
    _logger.note(f'Result = {est.eqtree}, RMSE = {rmse:.4f}, R2 = {r2:.4f}')

    result = {
        'date': time.strftime('%Y-%m-%d %H:%M:%S'),
        'host': gethostname(),
        'name': args.name,
        'success': str(rmse < 1e-6),
        'n_iter': len(est.records),
        'duration': est.records[-1]['time'],
        'exp': args.function,
        'result': str(est.eqtree),
        'rmse': rmse,
        'r2': r2,
        'sample_num': args.sample_num,
        'seed': args.seed,
        'status': status,
    }
    json.dump(result, open(os.path.join(args.save_dir, 'result.json'), 'w'), indent=4)

    try:
        df_list = []
        for node in est.pareto_front:
            df_list.append({'R2': node.r2, 'Complexity': node.complexity, 'Equation': node.phi})
        df_pareto = pd.DataFrame(df_list).sort_values(['R2', 'Complexity'], ascending=[False, True])
        header = f'{'R2':8}  {'Complexity':10}  {'Equation'}'
        lines = '\n'.join(f'{r2:8.4f}  {complexity:10d}  {phi}' for r2, complexity, phi in df_pareto.itertuples(index=False))
        _logger.info(f'Pareto front:\n{header}\n{lines}')
        os.makedirs(args.save_dir, exist_ok=True)
        save_path = os.path.join(args.save_dir, 'pareto_front.csv')
        df_pareto.to_csv(save_path, index=False)
        _logger.info(f'Pareto front saved to {save_path}')
    except Exception as e:
        _logger.error(
            f'Failed to save pareto front: '
            f'[{type(e).__name__}] {e}\n'
            f'{traceback.format_exc()}'
        )


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-f', '--function', type=str, default='f=x1+x2*sin(x3)', help='`f=...\' or `Feynman_xxx\'')
    parser.add_argument('-n', '--name', type=str, default=None)
    parser.add_argument('-s', '--seed', type=int, default=None)
    parser.add_argument('--method', type=str, choices=['eic-mcts', 'eic-gp'], default='eic-mcts')
    parser.add_argument('--sample_num', type=int, default=200)
    parser.add_argument('--max_len', type=int, default=30)
    parser.add_argument('--n_iter', type=int, default=10000)
    parser.add_argument('--max_var', type=int, default=10)
    parser.add_argument('--quiet', action='store_true')
    parser.add_argument('--keep_vars', action='store_true')
    parser.add_argument('--normalize_y', action='store_true')
    parser.add_argument('--normalize_X', action='store_true')
    parser.add_argument('--remove_abnormal', action='store_true')
    parser.add_argument('--cheat', action='store_true')
    parser.add_argument('--keep_name', action='store_true')
    parser.add_argument('--log_per_sec', type=float, default=5)
    parser.add_argument('--save_dir', type=str, default='./logs')
    parser.add_argument('--skip_existing', action='store_true')
    parser.add_argument('--time_limit', type=float, default=900)
    parser.add_argument('--no_time_limit', action='store_const', const=None, dest='time_limit')
    parser.add_argument('--c', type=float, default=1.414)
    parser.add_argument('--eta', type=float, default=0.999)
    parser.add_argument('--alpha', type=float, default=0.0)
    parser.add_argument('--target_noise', type=float, default=0.0)
    parser.add_argument('--ratio', type=float, default=1.0, help='set to 0.25 for blackbox experiments')
    args, unknown = parser.parse_known_args()

    if unknown:
        _logger.warning(f'unknown args: {unknown}')

    if not args.keep_name:
        def sanitize_filename(s: str) -> str:
            _illegal = re.compile(r'[<>:"/\\|?*\x00-\x1f]')  # Windows/一般不允许的字符
            s = s or ''
            s = s.strip()
            s = _illegal.sub('_', s)
            s = s.replace(' ', '_')
            return s or 'unnamed'

        date = datetime.datetime.now()
        hostname = socket.gethostname()
        yymmdd = date.strftime('%y%m%d')
        hhmmss = date.strftime('%H%M%S')
        safe_name = sanitize_filename(args.name)
        safe_host = sanitize_filename(hostname)
        args.name =  f"{yymmdd}_{safe_name}_{hhmmss}_{safe_host}"

        if args.skip_existing and any(Path(args.save_dir).glob(f'*_{safe_name}_*_*')):
            path = next(Path(args.save_dir).glob(f'*_{safe_name}_*_*'))
            _logger.warning(f'Existing experiment found for name={safe_name} in {path}, skip it.')
            exit(0)

    args.save_dir = os.path.join(args.save_dir, args.name)
    init_logger('src', args.name, Path(args.save_dir) / 'info.log')
    _logger.info(args)
    with open(Path(args.save_dir) / 'args.json', 'w', encoding='utf-8') as f:
        json.dump(vars(args), f, indent=4)

    if args.seed is None:
        args.seed = random.randint(0, 10000)
    args.function = args.function.replace(' ', '')

    setproctitle(f'{args.name}@YuZihan')
    def handler(signum, frame): raise KeyboardInterrupt
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)

    seed_all(args.seed)
    search(args)