[ { "path": ".idea/.gitignore", "content": "# Default ignored files\n/shelf/\n/workspace.xml\n# Editor-based HTTP Client requests\n/httpRequests/\n# Datasource local storage ignored files\n/dataSources/\n/dataSources.local.xml\n" }, { "path": ".idea/StockMixer.iml", "content": "\n\n \n \n \n \n \n \n \n" }, { "path": ".idea/inspectionProfiles/profiles_settings.xml", "content": "\n \n \n" }, { "path": ".idea/misc.xml", "content": "\n\n \n \n \n" }, { "path": ".idea/modules.xml", "content": "\n\n \n \n \n \n \n" }, { "path": ".idea/vcs.xml", "content": "\n\n \n \n \n \n" }, { "path": "README.md", "content": "# StockMixer\n\nOfficial code implementation and supplementary material of AAAI 2024 paper \"**StockMixer: A Simple yet Strong MLP-based Architecture for Stock Price Forecasting**\". This work proposes a lightweight and effective MLP-based architecture for stock price forecasting named StockMixer. It consists of indicator mixing, temporal mixing and stock mixing to capture complex correlations in the stock data. The end-to-end training flow of StockMixer is presented as follows:\n\n\n\n## Environment\n\n- Python 3.7\n- torch~=1.10.1\n- numpy~=1.21.5\n- PyYAML, pandas, tqdm, matplotlib\n\n## Dataset and Preprocessing\n\nThe original datasets(NASDAQ, NYSE and S&P500) are respectively available:\n\nNASDAQ/NYSE: [Temporal Relational Ranking for Stock Prediction](https://github.com/fulifeng/Temporal_Relational_Stock_Ranking)\n\nS&P500: [Efficient Integration of Multi-Order Dynamics and Internal Dynamics in Stock Movement Prediction](https://github.com/thanhtrunghuynh93/estimate)\n\nIn order to improve file reading speed, we process the raw data to generate corresponding .pkl or .npy files. Datasets are provided in the `dataset` folder. Because StockMixer does not require prior knowledge similar to graphs or hypergraphs, our preprocessed dataset did not provide either. You can find them from the original datasets.\n\n## Running the Code\n\n```\n# edit configurations in train.py\npython src/train.py\n```\n\n" }, { "path": "dataset/NYSE/eod_data.pkl", "content": "" }, { "path": "dataset/NYSE/gt_data.pkl", "content": "" }, { "path": "dataset/NYSE/mask_data.pkl", "content": "" }, { "path": "dataset/NYSE/price_data.pkl", "content": "" }, { "path": "dataset/SP500/sp500_ticker.csv", "content": "Symbol,Name,Sector\nMMM,3M,Industrials\nAOS,A. O. Smith,Industrials\nABT,Abbott Laboratories,Health Care\nABBV,AbbVie,Health Care\nABMD,Abiomed,Health Care\nACN,Accenture,Information Technology\nATVI,Activision Blizzard,Communication Services\nADM,ADM,Consumer Staples\nADBE,Adobe,Information Technology\nAAP,Advance Auto Parts,Consumer Discretionary\nAMD,Advanced Micro Devices,Information Technology\nAES,AES Corp,Utilities\nAFL,Aflac,Financials\nA,Agilent Technologies,Health Care\nAPD,Air Products & Chemicals,Materials\nAKAM,Akamai Technologies,Information Technology\nALK,Alaska Air Group,Industrials\nALB,Albemarle Corporation,Materials\nARE,Alexandria Real Estate Equities,Real Estate\nALGN,Align Technology,Health Care\nALLE,Allegion,Industrials\nLNT,Alliant Energy,Utilities\nALL,Allstate Corp,Financials\nGOOGL,Alphabet (Class A),Communication Services\nGOOG,Alphabet (Class C),Communication Services\nMO,Altria Group,Consumer Staples\nAMZN,Amazon,Consumer Discretionary\nAMCR,Amcor,Materials\nAEE,Ameren Corp,Utilities\nAAL,American Airlines Group,Industrials\nAEP,American Electric Power,Utilities\nAXP,American Express,Financials\nAIG,American International Group,Financials\nAMT,American Tower,Real Estate\nAWK,American Water Works,Utilities\nAMP,Ameriprise Financial,Financials\nABC,AmerisourceBergen,Health Care\nAME,Ametek,Industrials\nAMGN,Amgen,Health Care\nAPH,Amphenol,Information Technology\nADI,Analog Devices,Information Technology\nANSS,Ansys,Information Technology\nANTM,Anthem,Health Care\nAON,Aon,Financials\nAPA,APA Corporation,Energy\nAAPL,Apple,Information Technology\nAMAT,Applied Materials,Information Technology\nAPTV,Aptiv,Consumer Discretionary\nANET,Arista Networks,Information Technology\nAJG,Arthur J. Gallagher & Co.,Financials\nAIZ,Assurant,Financials\nT,AT&T,Communication Services\nATO,Atmos Energy,Utilities\nADSK,Autodesk,Information Technology\nADP,Automatic Data Processing,Information Technology\nAZO,AutoZone,Consumer Discretionary\nAVB,AvalonBay Communities,Real Estate\nAVY,Avery Dennison,Materials\nBKR,Baker Hughes,Energy\nBLL,Ball Corp,Materials\nBAC,Bank of America,Financials\nBBWI,Bath & Body Works Inc.,Consumer Discretionary\nBAX,Baxter International,Health Care\nBDX,Becton Dickinson,Health Care\nBRK-B,Berkshire Hathaway,Financials\nBBY,Best Buy,Consumer Discretionary\nBIO,Bio-Rad Laboratories,Health Care\nTECH,Bio-Techne,Health Care\nBIIB,Biogen,Health Care\nBLK,BlackRock,Financials\nBK,BNY Mellon,Financials\nBA,Boeing,Industrials\nBKNG,Booking Holdings,Consumer Discretionary\nBWA,BorgWarner,Consumer Discretionary\nBXP,Boston Properties,Real Estate\nBSX,Boston Scientific,Health Care\nBMY,Bristol Myers Squibb,Health Care\nAVGO,Broadcom,Information Technology\nBR,Broadridge Financial Solutions,Information Technology\nBRO,Brown & Brown,Financials\nBF-B,Brown–Forman,Consumer Staples\nCHRW,C. H. Robinson,Industrials\nCDNS,Cadence Design Systems,Information Technology\nCZR,Caesars Entertainment,Consumer Discretionary\nCPB,Campbell Soup,Consumer Staples\nCOF,Capital One Financial,Financials\nCAH,Cardinal Health,Health Care\nKMX,CarMax,Consumer Discretionary\nCCL,Carnival Corporation,Consumer Discretionary\nCARR,Carrier Global,Industrials\nCTLT,Catalent,Health Care\nCAT,Caterpillar,Industrials\nCBOE,Cboe Global Markets,Financials\nCBRE,CBRE,Real Estate\nCDW,CDW,Information Technology\nCE,Celanese,Materials\nCNC,Centene Corporation,Health Care\nCNP,CenterPoint Energy,Utilities\nCDAY,Ceridian,Information Technology\nCERN,Cerner,Health Care\nCF,CF Industries,Materials\nCRL,Charles River Laboratories,Health Care\nSCHW,Charles Schwab Corporation,Financials\nCHTR,Charter Communications,Communication Services\nCVX,Chevron Corporation,Energy\nCMG,Chipotle Mexican Grill,Consumer Discretionary\nCB,Chubb,Financials\nCHD,Church & Dwight,Consumer Staples\nCI,Cigna,Health Care\nCINF,Cincinnati Financial,Financials\nCTAS,Cintas Corporation,Industrials\nCSCO,Cisco Systems,Information Technology\nC,Citigroup,Financials\nCFG,Citizens Financial Group,Financials\nCTXS,Citrix Systems,Information Technology\nCLX,Clorox,Consumer Staples\nCME,CME Group,Financials\nCMS,CMS Energy,Utilities\nKO,Coca-Cola Company,Consumer Staples\nCTSH,Cognizant Technology Solutions,Information Technology\nCL,Colgate-Palmolive,Consumer Staples\nCMCSA,Comcast,Communication Services\nCMA,Comerica,Financials\nCAG,Conagra Brands,Consumer Staples\nCOP,ConocoPhillips,Energy\nED,Consolidated Edison,Utilities\nSTZ,Constellation Brands,Consumer Staples\nCPRT,Copart,Industrials\nGLW,Corning,Information Technology\nCTVA,Corteva,Materials\nCOST,Costco,Consumer Staples\nCTRA,Coterra,Energy\nCCI,Crown Castle,Real Estate\nCSX,CSX,Industrials\nCMI,Cummins,Industrials\nCVS,CVS Health,Health Care\nDHI,D. R. Horton,Consumer Discretionary\nDHR,Danaher Corporation,Health Care\nDRI,Darden Restaurants,Consumer Discretionary\nDVA,DaVita,Health Care\nDE,Deere & Co.,Industrials\nDAL,Delta Air Lines,Industrials\nXRAY,Dentsply Sirona,Health Care\nDVN,Devon Energy,Energy\nDXCM,DexCom,Health Care\nFANG,Diamondback Energy,Energy\nDLR,Digital Realty Trust,Real Estate\nDFS,Discover Financial Services,Financials\nDISCA,Discovery (Series A),Communication Services\nDISCK,Discovery (Series C),Communication Services\nDISH,Dish Network,Communication Services\nDG,Dollar General,Consumer Discretionary\nDLTR,Dollar Tree,Consumer Discretionary\nD,Dominion Energy,Utilities\nDPZ,Domino's Pizza,Consumer Discretionary\nDOV,Dover Corporation,Industrials\nDOW,Dow,Materials\nDTE,DTE Energy,Utilities\nDUK,Duke Energy,Utilities\nDRE,Duke Realty Corp,Real Estate\nDD,DuPont,Materials\nDXC,DXC Technology,Information Technology\nEMN,Eastman Chemical,Materials\nETN,Eaton Corporation,Industrials\nEBAY,eBay,Consumer Discretionary\nECL,Ecolab,Materials\nEIX,Edison International,Utilities\nEW,Edwards Lifesciences,Health Care\nEA,Electronic Arts,Communication Services\nLLY,Eli Lilly & Co,Health Care\nEMR,Emerson Electric Company,Industrials\nENPH,Enphase Energy,Information Technology\nETR,Entergy,Utilities\nEOG,EOG Resources,Energy\nEFX,Equifax,Industrials\nEQIX,Equinix,Real Estate\nEQR,Equity Residential,Real Estate\nESS,Essex Property Trust,Real Estate\nEL,Estée Lauder Companies,Consumer Staples\nETSY,Etsy,Consumer Discretionary\nRE,Everest Re,Financials\nEVRG,Evergy,Utilities\nES,Eversource Energy,Utilities\nEXC,Exelon,Utilities\nEXPE,Expedia Group,Consumer Discretionary\nEXPD,Expeditors,Industrials\nEXR,Extra Space Storage,Real Estate\nXOM,ExxonMobil,Energy\nFFIV,F5 Networks,Information Technology\nFB,Facebook,Communication Services\nFAST,Fastenal,Industrials\nFRT,Federal Realty Investment Trust,Real Estate\nFDX,FedEx,Industrials\nFIS,Fidelity National Information Services,Information Technology\nFITB,Fifth Third Bancorp,Financials\nFRC,First Republic Bank,Financials\nFE,FirstEnergy,Utilities\nFISV,Fiserv,Information Technology\nFLT,Fleetcor,Information Technology\nFMC,FMC Corporation,Materials\nF,Ford,Consumer Discretionary\nFTNT,Fortinet,Information Technology\nFTV,Fortive,Industrials\nFBHS,Fortune Brands Home & Security,Industrials\nFOXA,Fox Corporation (Class A),Communication Services\nFOX,Fox Corporation (Class B),Communication Services\nBEN,Franklin Resources,Financials\nFCX,Freeport-McMoRan,Materials\nGPS,Gap,Consumer Discretionary\nGRMN,Garmin,Consumer Discretionary\nIT,Gartner,Information Technology\nGNRC,Generac Holdings,Industrials\nGD,General Dynamics,Industrials\nGE,General Electric,Industrials\nGIS,General Mills,Consumer Staples\nGM,General Motors,Consumer Discretionary\nGPC,Genuine Parts,Consumer Discretionary\nGILD,Gilead Sciences,Health Care\nGPN,Global Payments,Information Technology\nGL,Globe Life,Financials\nGS,Goldman Sachs,Financials\nHAL,Halliburton,Energy\nHBI,Hanesbrands,Consumer Discretionary\nHAS,Hasbro,Consumer Discretionary\nHCA,HCA Healthcare,Health Care\nPEAK,Healthpeak Properties,Real Estate\nHSIC,Henry Schein,Health Care\nHES,Hess Corporation,Energy\nHPE,Hewlett Packard Enterprise,Information Technology\nHLT,Hilton Worldwide,Consumer Discretionary\nHOLX,Hologic,Health Care\nHD,Home Depot,Consumer Discretionary\nHON,Honeywell,Industrials\nHRL,Hormel,Consumer Staples\nHST,Host Hotels & Resorts,Real Estate\nHWM,Howmet Aerospace,Industrials\nHPQ,HP,Information Technology\nHUM,Humana,Health Care\nHBAN,Huntington Bancshares,Financials\nHII,Huntington Ingalls Industries,Industrials\nIBM,IBM,Information Technology\nIEX,IDEX Corporation,Industrials\nIDXX,Idexx Laboratories,Health Care\nINFO,IHS Markit,Industrials\nITW,Illinois Tool Works,Industrials\nILMN,Illumina,Health Care\nINCY,Incyte,Health Care\nIR,Ingersoll Rand,Industrials\nINTC,Intel,Information Technology\nICE,Intercontinental Exchange,Financials\nIFF,International Flavors & Fragrances,Materials\nIP,International Paper,Materials\nIPG,Interpublic Group,Communication Services\nINTU,Intuit,Information Technology\nISRG,Intuitive Surgical,Health Care\nIVZ,Invesco,Financials\nIPGP,IPG Photonics,Information Technology\nIQV,IQVIA,Health Care\nIRM,Iron Mountain,Real Estate\nJBHT,J. B. Hunt,Industrials\nJKHY,Jack Henry & Associates,Information Technology\nJ,Jacobs Engineering Group,Industrials\nSJM,JM Smucker,Consumer Staples\nJNJ,Johnson & Johnson,Health Care\nJCI,Johnson Controls,Industrials\nJPM,JPMorgan Chase,Financials\nJNPR,Juniper Networks,Information Technology\nKSU,Kansas City Southern,Industrials\nK,Kellogg's,Consumer Staples\nKEY,KeyCorp,Financials\nKEYS,Keysight Technologies,Information Technology\nKMB,Kimberly-Clark,Consumer Staples\nKIM,Kimco Realty,Real Estate\nKMI,Kinder Morgan,Energy\nKLAC,KLA Corporation,Information Technology\nKHC,Kraft Heinz,Consumer Staples\nKR,Kroger,Consumer Staples\nLHX,L3Harris Technologies,Industrials\nLH,LabCorp,Health Care\nLRCX,Lam Research,Information Technology\nLW,Lamb Weston,Consumer Staples\nLVS,Las Vegas Sands,Consumer Discretionary\nLEG,Leggett & Platt,Consumer Discretionary\nLDOS,Leidos,Industrials\nLEN,Lennar,Consumer Discretionary\nLNC,Lincoln National,Financials\nLIN,Linde,Materials\nLYV,Live Nation Entertainment,Communication Services\nLKQ,LKQ Corporation,Consumer Discretionary\nLMT,Lockheed Martin,Industrials\nL,Loews Corporation,Financials\nLOW,Lowe's,Consumer Discretionary\nLUMN,Lumen Technologies,Communication Services\nLYB,LyondellBasell,Materials\nMTB,M&T Bank,Financials\nMRO,Marathon Oil,Energy\nMPC,Marathon Petroleum,Energy\nMKTX,MarketAxess,Financials\nMAR,Marriott International,Consumer Discretionary\nMMC,Marsh & McLennan,Financials\nMLM,Martin Marietta Materials,Materials\nMAS,Masco,Industrials\nMA,Mastercard,Information Technology\nMTCH,Match Group,Communication Services\nMKC,McCormick & Company,Consumer Staples\nMCD,McDonald's,Consumer Discretionary\nMCK,McKesson Corporation,Health Care\nMDT,Medtronic,Health Care\nMRK,Merck & Co.,Health Care\nMET,MetLife,Financials\nMTD,Mettler Toledo,Health Care\nMGM,MGM Resorts International,Consumer Discretionary\nMCHP,Microchip Technology,Information Technology\nMU,Micron Technology,Information Technology\nMSFT,Microsoft,Information Technology\nMAA,Mid-America Apartments,Real Estate\nMRNA,Moderna,Health Care\nMHK,Mohawk Industries,Consumer Discretionary\nTAP,Molson Coors Beverage Company,Consumer Staples\nMDLZ,Mondelez International,Consumer Staples\nMPWR,Monolithic Power Systems,Information Technology\nMNST,Monster Beverage,Consumer Staples\nMCO,Moody's Corporation,Financials\nMS,Morgan Stanley,Financials\nMSI,Motorola Solutions,Information Technology\nMSCI,MSCI,Financials\nNDAQ,Nasdaq,Financials\nNTAP,NetApp,Information Technology\nNFLX,Netflix,Communication Services\nNWL,Newell Brands,Consumer Discretionary\nNEM,Newmont,Materials\nNWSA,News Corp (Class A),Communication Services\nNWS,News Corp (Class B),Communication Services\nNEE,NextEra Energy,Utilities\nNLSN,Nielsen Holdings,Industrials\nNKE,Nike,Consumer Discretionary\nNI,NiSource,Utilities\nNSC,Norfolk Southern,Industrials\nNTRS,Northern Trust,Financials\nNOC,Northrop Grumman,Industrials\nNLOK,NortonLifeLock,Information Technology\nNCLH,Norwegian Cruise Line Holdings,Consumer Discretionary\nNRG,NRG Energy,Utilities\nNUE,Nucor,Materials\nNVDA,Nvidia,Information Technology\nNVR,NVR,Consumer Discretionary\nNXPI,NXP,Information Technology\nORLY,O'Reilly Automotive,Consumer Discretionary\nOXY,Occidental Petroleum,Energy\nODFL,Old Dominion Freight Line,Industrials\nOMC,Omnicom Group,Communication Services\nOKE,Oneok,Energy\nORCL,Oracle,Information Technology\nOGN,Organon & Co.,Health Care\nOTIS,Otis Worldwide,Industrials\nPCAR,Paccar,Industrials\nPKG,Packaging Corporation of America,Materials\nPH,Parker-Hannifin,Industrials\nPAYX,Paychex,Information Technology\nPAYC,Paycom,Information Technology\nPYPL,PayPal,Information Technology\nPENN,Penn National Gaming,Consumer Discretionary\nPNR,Pentair,Industrials\nPBCT,People's United Financial,Financials\nPEP,PepsiCo,Consumer Staples\nPKI,PerkinElmer,Health Care\nPFE,Pfizer,Health Care\nPM,Philip Morris International,Consumer Staples\nPSX,Phillips 66,Energy\nPNW,Pinnacle West Capital,Utilities\nPXD,Pioneer Natural Resources,Energy\nPNC,PNC Financial Services,Financials\nPOOL,Pool Corporation,Consumer Discretionary\nPPG,PPG Industries,Materials\nPPL,PPL,Utilities\nPFG,Principal Financial Group,Financials\nPG,Procter & Gamble,Consumer Staples\nPGR,Progressive Corporation,Financials\nPLD,Prologis,Real Estate\nPRU,Prudential Financial,Financials\nPTC,PTC,Information Technology\nPEG,Public Service Enterprise Group,Utilities\nPSA,Public Storage,Real Estate\nPHM,PulteGroup,Consumer Discretionary\nPVH,PVH,Consumer Discretionary\nQRVO,Qorvo,Information Technology\nQCOM,Qualcomm,Information Technology\nPWR,Quanta Services,Industrials\nDGX,Quest Diagnostics,Health Care\nRL,Ralph Lauren Corporation,Consumer Discretionary\nRJF,Raymond James Financial,Financials\nRTX,Raytheon Technologies,Industrials\nO,Realty Income Corporation,Real Estate\nREG,Regency Centers,Real Estate\nREGN,Regeneron Pharmaceuticals,Health Care\nRF,Regions Financial Corporation,Financials\nRSG,Republic Services,Industrials\nRMD,ResMed,Health Care\nRHI,Robert Half International,Industrials\nROK,Rockwell Automation,Industrials\nROL,Rollins,Industrials\nROP,Roper Technologies,Industrials\nROST,Ross Stores,Consumer Discretionary\nRCL,Royal Caribbean Group,Consumer Discretionary\nSPGI,S&P Global,Financials\nCRM,Salesforce,Information Technology\nSBAC,SBA Communications,Real Estate\nSLB,Schlumberger,Energy\nSTX,Seagate Technology,Information Technology\nSEE,Sealed Air,Materials\nSRE,Sempra Energy,Utilities\nNOW,ServiceNow,Information Technology\nSHW,Sherwin-Williams,Materials\nSPG,Simon Property Group,Real Estate\nSWKS,Skyworks Solutions,Information Technology\nSNA,Snap-on,Industrials\nSO,Southern Company,Utilities\nLUV,Southwest Airlines,Industrials\nSWK,Stanley Black & Decker,Industrials\nSBUX,Starbucks,Consumer Discretionary\nSTT,State Street Corporation,Financials\nSTE,Steris,Health Care\nSYK,Stryker Corporation,Health Care\nSIVB,SVB Financial,Financials\nSYF,Synchrony Financial,Financials\nSNPS,Synopsys,Information Technology\nSYY,Sysco,Consumer Staples\nTMUS,T-Mobile US,Communication Services\nTROW,T. Rowe Price,Financials\nTTWO,Take-Two Interactive,Communication Services\nTPR,Tapestry,Consumer Discretionary\nTGT,Target Corporation,Consumer Discretionary\nTEL,TE Connectivity,Information Technology\nTDY,Teledyne Technologies,Industrials\nTFX,Teleflex,Health Care\nTER,Teradyne,Information Technology\nTSLA,Tesla,Consumer Discretionary\nTXN,Texas Instruments,Information Technology\nTXT,Textron,Industrials\nCOO,The Cooper Companies,Health Care\nHIG,The Hartford,Financials\nHSY,The Hershey Company,Consumer Staples\nMOS,The Mosaic Company,Materials\nTRV,The Travelers Companies,Financials\nDIS,The Walt Disney Company,Communication Services\nTMO,Thermo Fisher Scientific,Health Care\nTJX,TJX Companies,Consumer Discretionary\nTSCO,Tractor Supply Company,Consumer Discretionary\nTT,Trane Technologies,Industrials\nTDG,TransDigm Group,Industrials\nTRMB,Trimble,Information Technology\nTFC,Truist Financial,Financials\nTWTR,Twitter,Communication Services\nTYL,Tyler Technologies,Information Technology\nTSN,Tyson Foods,Consumer Staples\nUSB,U.S. Bancorp,Financials\nUDR,UDR,Real Estate\nULTA,Ulta Beauty,Consumer Discretionary\nUAA,Under Armour (Class A),Consumer Discretionary\nUA,Under Armour (Class C),Consumer Discretionary\nUNP,Union Pacific,Industrials\nUAL,United Airlines,Industrials\nUPS,United Parcel Service,Industrials\nURI,United Rentals,Industrials\nUNH,UnitedHealth Group,Health Care\nUHS,Universal Health Services,Health Care\nVLO,Valero Energy,Energy\nVTR,Ventas,Real Estate\nVRSN,Verisign,Information Technology\nVRSK,Verisk Analytics,Industrials\nVZ,Verizon Communications,Communication Services\nVRTX,Vertex Pharmaceuticals,Health Care\nVFC,VF Corporation,Consumer Discretionary\nVIAC,ViacomCBS,Communication Services\nVTRS,Viatris,Health Care\nV,Visa,Information Technology\nVNO,Vornado Realty Trust,Real Estate\nVMC,Vulcan Materials,Materials\nWRB,W. R. Berkley Corporation,Financials\nGWW,W. W. Grainger,Industrials\nWAB,Wabtec,Industrials\nWBA,Walgreens Boots Alliance,Consumer Staples\nWMT,Walmart,Consumer Staples\nWM,Waste Management,Industrials\nWAT,Waters Corporation,Health Care\nWEC,WEC Energy Group,Utilities\nWFC,Wells Fargo,Financials\nWELL,Welltower,Real Estate\nWST,West Pharmaceutical Services,Health Care\nWDC,Western Digital,Information Technology\nWU,Western Union,Information Technology\nWRK,WestRock,Materials\nWY,Weyerhaeuser,Real Estate\nWHR,Whirlpool Corporation,Consumer Discretionary\nWMB,Williams Companies,Energy\nWLTW,Willis Towers Watson,Financials\nWYNN,Wynn Resorts,Consumer Discretionary\nXEL,Xcel Energy,Utilities\nXLNX,Xilinx,Information Technology\nXYL,Xylem,Industrials\nYUM,Yum! Brands,Consumer Discretionary\nZBRA,Zebra Technologies,Information Technology\nZBH,Zimmer Biomet,Health Care\nZION,Zions Bancorp,Financials\nZTS,Zoetis,Health Care" }, { "path": "requirements.txt", "content": "tqdm==4.64.1\npandas==2.0.3\npyyaml==6.0\nnumpy==1.22.1\nmatplotlib==3.5.1" }, { "path": "src/evaluator.py", "content": "import numpy as np\nimport pandas as pd\n\n\ndef evaluate(prediction, ground_truth, mask, report=False):\n assert ground_truth.shape == prediction.shape, 'shape mis-match'\n performance = {}\n # mse\n performance['mse'] = np.linalg.norm((prediction - ground_truth) * mask) ** 2 / np.sum(mask)\n # IC\n df_pred = pd.DataFrame(prediction * mask)\n df_gt = pd.DataFrame(ground_truth * mask)\n\n ic = []\n mrr_top = 0.0\n all_miss_days_top = 0\n bt_long = 1.0\n bt_long5 = 1.0\n bt_long10 = 1.0\n irr = 0.0\n sharpe_li5 = []\n prec_10 = []\n\n for i in range(prediction.shape[1]):\n # IC\n ic.append(df_pred[i].corr(df_gt[i]))\n\n rank_gt = np.argsort(ground_truth[:, i])\n gt_top1 = set()\n gt_top5 = set()\n gt_top10 = set()\n\n for j in range(1, prediction.shape[0] + 1):\n cur_rank = rank_gt[-1 * j]\n if mask[cur_rank][i] < 0.5:\n continue\n if len(gt_top1) < 1:\n gt_top1.add(cur_rank)\n if len(gt_top5) < 5:\n gt_top5.add(cur_rank)\n if len(gt_top10) < 10:\n gt_top10.add(cur_rank)\n\n rank_pre = np.argsort(prediction[:, i])\n pre_top1 = set()\n pre_top5 = set()\n pre_top10 = set()\n for j in range(1, prediction.shape[0] + 1):\n cur_rank = rank_pre[-1 * j]\n if mask[cur_rank][i] < 0.5:\n continue\n if len(pre_top1) < 1:\n pre_top1.add(cur_rank)\n if len(pre_top5) < 5:\n pre_top5.add(cur_rank)\n if len(pre_top10) < 10:\n pre_top10.add(cur_rank)\n\n top1_pos_in_gt = 0\n for j in range(1, prediction.shape[0] + 1):\n cur_rank = rank_gt[-1 * j]\n if mask[cur_rank][i] < 0.5:\n continue\n else:\n top1_pos_in_gt += 1\n if cur_rank in pre_top1:\n break\n if top1_pos_in_gt == 0:\n all_miss_days_top += 1\n else:\n mrr_top += 1.0 / top1_pos_in_gt\n\n real_ret_rat_top = ground_truth[list(pre_top1)[0]][i]\n bt_long += real_ret_rat_top\n gt_irr = 0.0\n\n for gt in gt_top10:\n gt_irr += ground_truth[gt][i]\n\n real_ret_rat_top5 = 0\n for pre in pre_top5:\n real_ret_rat_top5 += ground_truth[pre][i]\n irr += real_ret_rat_top5\n real_ret_rat_top5 /= 5\n bt_long5 += real_ret_rat_top5\n\n prec = 0.0\n real_ret_rat_top10 = 0\n for pre in pre_top10:\n real_ret_rat_top10 += ground_truth[pre][i]\n prec += (ground_truth[pre][i] >= 0)\n prec_10.append(prec / 10)\n real_ret_rat_top10 /= 10\n bt_long10 += real_ret_rat_top10\n sharpe_li5.append(real_ret_rat_top5)\n\n performance['IC'] = np.mean(ic)\n performance['RIC'] = np.mean(ic) / np.std(ic)\n sharpe_li5 = np.array(sharpe_li5)\n performance['sharpe5'] = (np.mean(sharpe_li5)/np.std(sharpe_li5))*15.87\n performance['prec_10'] = np.mean(prec_10)\n return performance\n\n\n\n" }, { "path": "src/load_data.py", "content": "import numpy as np\nimport os\nfrom tqdm import tqdm\n\n\ndef load_EOD_data(data_path, market_name, tickers, steps=1):\n eod_data = []\n masks = []\n ground_truth = []\n base_price = []\n for index, ticker in enumerate(tqdm(tickers)):\n single_EOD = np.genfromtxt(\n os.path.join(data_path, market_name + '_' + ticker + '_1.csv'),\n dtype=np.float32, delimiter=',', skip_header=False\n )\n if market_name == 'NASDAQ':\n # remove the last day since lots of missing data\n single_EOD = single_EOD[:-1, :]\n if index == 0:\n print('single EOD data shape:', single_EOD.shape)\n eod_data = np.zeros([len(tickers), single_EOD.shape[0],\n single_EOD.shape[1] - 1], dtype=np.float32)\n masks = np.ones([len(tickers), single_EOD.shape[0]],\n dtype=np.float32)\n ground_truth = np.zeros([len(tickers), single_EOD.shape[0]],\n dtype=np.float32)\n base_price = np.zeros([len(tickers), single_EOD.shape[0]],\n dtype=np.float32)\n for row in range(single_EOD.shape[0]):\n if abs(single_EOD[row][-1] + 1234) < 1e-8:\n masks[index][row] = 0.0\n elif row > steps - 1 and abs(single_EOD[row - steps][-1] + 1234) \\\n > 1e-8:\n ground_truth[index][row] = \\\n (single_EOD[row][-1] - single_EOD[row - steps][-1]) / \\\n single_EOD[row - steps][-1]\n for col in range(single_EOD.shape[1]):\n if abs(single_EOD[row][col] + 1234) < 1e-8:\n single_EOD[row][col] = 1.1\n eod_data[index, :, :] = single_EOD[:, 1:]\n base_price[index, :] = single_EOD[:, -1]\n return eod_data, masks, ground_truth, base_price\n\n\ndef load_graph_relation_data(relation_file, lap=False):\n relation_encoding = np.load(relation_file)\n print('relation encoding shape:', relation_encoding.shape)\n rel_shape = [relation_encoding.shape[0], relation_encoding.shape[1]]\n mask_flags = np.equal(np.zeros(rel_shape, dtype=int),\n np.sum(relation_encoding, axis=2))\n ajacent = np.where(mask_flags, np.zeros(rel_shape, dtype=float),\n np.ones(rel_shape, dtype=float))\n degree = np.sum(ajacent, axis=0)\n for i in range(len(degree)):\n degree[i] = 1.0 / degree[i]\n np.sqrt(degree, degree)\n deg_neg_half_power = np.diag(degree)\n if lap:\n return np.identity(ajacent.shape[0], dtype=float) - np.dot(\n np.dot(deg_neg_half_power, ajacent), deg_neg_half_power)\n else:\n return np.dot(np.dot(deg_neg_half_power, ajacent), deg_neg_half_power)\n\n\ndef load_relation_data(relation_file):\n relation_encoding = np.load(relation_file)\n rel_shape = [relation_encoding.shape[0], relation_encoding.shape[1]]\n mask_flags = np.equal(np.zeros(rel_shape, dtype=int),\n np.sum(relation_encoding, axis=2))\n mask = np.where(mask_flags, np.ones(rel_shape) * -1e9, np.zeros(rel_shape))\n return relation_encoding, mask\n\n\ndef build_SFM_data(data_path, market_name, tickers):\n eod_data = []\n for index, ticker in enumerate(tickers):\n single_EOD = np.genfromtxt(\n os.path.join(data_path, market_name + '_' + ticker + '_1.csv'),\n dtype=np.float32, delimiter=',', skip_header=False\n )\n if index == 0:\n print('single EOD data shape:', single_EOD.shape)\n eod_data = np.zeros([len(tickers), single_EOD.shape[0]],\n dtype=np.float32)\n\n for row in range(single_EOD.shape[0]):\n if abs(single_EOD[row][-1] + 1234) < 1e-8:\n if row < 3:\n for i in range(row + 1, single_EOD.shape[0]):\n if abs(single_EOD[i][-1] + 1234) > 1e-8:\n eod_data[index][row] = single_EOD[i][-1]\n break\n else:\n eod_data[index][row] = np.sum(\n eod_data[index, row - 3:row]) / 3\n else:\n eod_data[index][row] = single_EOD[row][-1]\n np.save(market_name + '_sfm_data', eod_data)" }, { "path": "src/model.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nacv = nn.GELU()\n\ndef get_loss(prediction, ground_truth, base_price, mask, batch_size, alpha):\n device = prediction.device\n all_one = torch.ones(batch_size, 1, dtype=torch.float32).to(device)\n return_ratio = torch.div(torch.sub(prediction, base_price), base_price)\n reg_loss = F.mse_loss(return_ratio * mask, ground_truth * mask)\n pre_pw_dif = torch.sub(\n return_ratio @ all_one.t(),\n all_one @ return_ratio.t()\n )\n gt_pw_dif = torch.sub(\n all_one @ ground_truth.t(),\n ground_truth @ all_one.t()\n )\n mask_pw = mask @ mask.t()\n rank_loss = torch.mean(\n F.relu(pre_pw_dif * gt_pw_dif * mask_pw)\n )\n loss = reg_loss + alpha * rank_loss\n return loss, reg_loss, rank_loss, return_ratio\n\n\nclass MixerBlock(nn.Module):\n def __init__(self, mlp_dim, hidden_dim, dropout=0.0):\n super(MixerBlock, self).__init__()\n self.mlp_dim = mlp_dim\n self.dropout = dropout\n\n self.dense_1 = nn.Linear(mlp_dim, hidden_dim)\n self.LN = acv\n self.dense_2 = nn.Linear(hidden_dim, mlp_dim)\n\n def forward(self, x):\n x = self.dense_1(x)\n x = self.LN(x)\n if self.dropout != 0.0:\n x = F.dropout(x, p=self.dropout)\n x = self.dense_2(x)\n if self.dropout != 0.0:\n x = F.dropout(x, p=self.dropout)\n return x\n\n\nclass Mixer2d(nn.Module):\n def __init__(self, time_steps, channels):\n super(Mixer2d, self).__init__()\n self.LN_1 = nn.LayerNorm([time_steps, channels])\n self.LN_2 = nn.LayerNorm([time_steps, channels])\n self.timeMixer = MixerBlock(time_steps, time_steps)\n self.channelMixer = MixerBlock(channels, channels)\n\n def forward(self, inputs):\n x = self.LN_1(inputs)\n x = x.permute(0, 2, 1)\n x = self.timeMixer(x)\n x = x.permute(0, 2, 1)\n\n x = self.LN_2(x + inputs)\n y = self.channelMixer(x)\n return x + y\n\n\nclass TriU(nn.Module):\n def __init__(self, time_step):\n super(TriU, self).__init__()\n self.time_step = time_step\n self.triU = nn.ParameterList(\n [\n nn.Linear(i + 1, 1)\n for i in range(time_step)\n ]\n )\n\n def forward(self, inputs):\n x = self.triU[0](inputs[:, :, 0].unsqueeze(-1))\n for i in range(1, self.time_step):\n x = torch.cat([x, self.triU[i](inputs[:, :, 0:i + 1])], dim=-1)\n return x\n\n\nclass TimeMixerBlock(nn.Module):\n def __init__(self, time_step):\n super(TimeMixerBlock, self).__init__()\n self.time_step = time_step\n self.dense_1 = TriU(time_step)\n self.LN = acv\n self.dense_2 = TriU(time_step)\n\n def forward(self, x):\n x = self.dense_1(x)\n x = self.LN(x)\n x = self.dense_2(x)\n return x\n\n\nclass MultiScaleTimeMixer(nn.Module):\n def __init__(self, time_step, channel, scale_count=1):\n super(MultiScaleTimeMixer, self).__init__()\n self.time_step = time_step\n self.scale_count = scale_count\n self.mix_layer = nn.ParameterList([nn.Sequential(\n nn.Conv1d(in_channels=channel, out_channels=channel, kernel_size=2 ** i, stride=2 ** i),\n TriU(int(time_step / 2 ** i)),\n nn.Hardswish(),\n TriU(int(time_step / 2 ** i))\n ) for i in range(scale_count)])\n self.mix_layer[0] = nn.Sequential(\n nn.LayerNorm([time_step, channel]),\n TriU(int(time_step)),\n nn.Hardswish(),\n TriU(int(time_step))\n )\n\n def forward(self, x):\n x = x.permute(0, 2, 1)\n y = self.mix_layer[0](x)\n for i in range(1, self.scale_count):\n y = torch.cat((y, self.mix_layer[i](x)), dim=-1)\n return y\n\n\nclass Mixer2dTriU(nn.Module):\n def __init__(self, time_steps, channels):\n super(Mixer2dTriU, self).__init__()\n self.LN_1 = nn.LayerNorm([time_steps, channels])\n self.LN_2 = nn.LayerNorm([time_steps, channels])\n self.timeMixer = TriU(time_steps)\n self.channelMixer = MixerBlock(channels, channels)\n\n def forward(self, inputs):\n x = self.LN_1(inputs)\n x = x.permute(0, 2, 1)\n x = self.timeMixer(x)\n x = x.permute(0, 2, 1)\n\n x = self.LN_2(x + inputs)\n y = self.channelMixer(x)\n return x + y\n\n\nclass MultTime2dMixer(nn.Module):\n def __init__(self, time_step, channel, scale_dim=8):\n super(MultTime2dMixer, self).__init__()\n self.mix_layer = Mixer2dTriU(time_step, channel)\n self.scale_mix_layer = Mixer2dTriU(scale_dim, channel)\n\n def forward(self, inputs, y):\n y = self.scale_mix_layer(y)\n x = self.mix_layer(inputs)\n return torch.cat([inputs, x, y], dim=1)\n\n\nclass NoGraphMixer(nn.Module):\n def __init__(self, stocks, hidden_dim=20):\n super(NoGraphMixer, self).__init__()\n self.dense1 = nn.Linear(stocks, hidden_dim)\n self.activation = nn.Hardswish()\n self.dense2 = nn.Linear(hidden_dim, stocks)\n self.layer_norm_stock = nn.LayerNorm(stocks)\n\n def forward(self, inputs):\n x = inputs\n x = x.permute(1, 0)\n x = self.layer_norm_stock(x)\n x = self.dense1(x)\n x = self.activation(x)\n x = self.dense2(x)\n x = x.permute(1, 0)\n return x\n\n\nclass StockMixer(nn.Module):\n def __init__(self, stocks, time_steps, channels, market, scale):\n super(StockMixer, self).__init__()\n scale_dim = 8\n self.mixer = MultTime2dMixer(time_steps, channels, scale_dim=scale_dim)\n self.channel_fc = nn.Linear(channels, 1)\n self.time_fc = nn.Linear(time_steps * 2 + scale_dim, 1)\n self.conv = nn.Conv1d(in_channels=channels, out_channels=channels, kernel_size=2, stride=2)\n self.stock_mixer = NoGraphMixer(stocks, market)\n self.time_fc_ = nn.Linear(time_steps * 2 + scale_dim, 1)\n\n def forward(self, inputs):\n x = inputs.permute(0, 2, 1)\n x = self.conv(x)\n x = x.permute(0, 2, 1)\n y = self.mixer(inputs, x)\n y = self.channel_fc(y).squeeze(-1)\n\n z = self.stock_mixer(y)\n y = self.time_fc(y)\n z = self.time_fc_(z)\n return y + z\n\n" }, { "path": "src/train.py", "content": "import random\nimport numpy as np\nimport os\nimport torch as torch\nfrom load_data import load_EOD_data\nfrom evaluator import evaluate\nfrom model import get_loss, StockMixer\nimport pickle\n\n\nnp.random.seed(123456789)\ntorch.random.manual_seed(12345678)\ndevice = torch.device(\"cuda\") if torch.cuda.is_available() else 'cpu'\n\ndata_path = '../dataset'\nmarket_name = 'NASDAQ'\nrelation_name = 'wikidata'\nstock_num = 1026\nlookback_length = 16\nepochs = 100\nvalid_index = 756\ntest_index = 1008\nfea_num = 5\nmarket_num = 20\nsteps = 1\nlearning_rate = 0.001\nalpha = 0.1\nscale_factor = 3\nactivation = 'GELU'\n\ndataset_path = '../dataset/' + market_name\nif market_name == \"SP500\":\n data = np.load('../dataset/SP500/SP500.npy')\n data = data[:, 915:, :]\n price_data = data[:, :, -1]\n mask_data = np.ones((data.shape[0], data.shape[1]))\n eod_data = data\n gt_data = np.zeros((data.shape[0], data.shape[1]))\n for ticket in range(0, data.shape[0]):\n for row in range(1, data.shape[1]):\n gt_data[ticket][row] = (data[ticket][row][-1] - data[ticket][row - steps][-1]) / \\\n data[ticket][row - steps][-1]\nelse:\n with open(os.path.join(dataset_path, \"eod_data.pkl\"), \"rb\") as f:\n eod_data = pickle.load(f)\n with open(os.path.join(dataset_path, \"mask_data.pkl\"), \"rb\") as f:\n mask_data = pickle.load(f)\n with open(os.path.join(dataset_path, \"gt_data.pkl\"), \"rb\") as f:\n gt_data = pickle.load(f)\n with open(os.path.join(dataset_path, \"price_data.pkl\"), \"rb\") as f:\n price_data = pickle.load(f)\n\ntrade_dates = mask_data.shape[1]\nmodel = StockMixer(\n stocks=stock_num,\n time_steps=lookback_length,\n channels=fea_num,\n market=market_num,\n scale=scale_factor\n).to(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\nbest_valid_loss = np.inf\nbest_valid_perf = None\nbest_test_perf = None\nbatch_offsets = np.arange(start=0, stop=valid_index, dtype=int)\n\n\ndef validate(start_index, end_index):\n with torch.no_grad():\n cur_valid_pred = np.zeros([stock_num, end_index - start_index], dtype=float)\n cur_valid_gt = np.zeros([stock_num, end_index - start_index], dtype=float)\n cur_valid_mask = np.zeros([stock_num, end_index - start_index], dtype=float)\n loss = 0.\n reg_loss = 0.\n rank_loss = 0.\n for cur_offset in range(start_index - lookback_length - steps + 1, end_index - lookback_length - steps + 1):\n data_batch, mask_batch, price_batch, gt_batch = map(\n\n lambda x: torch.Tensor(x).to(device),\n get_batch(cur_offset)\n )\n prediction = model(data_batch)\n cur_loss, cur_reg_loss, cur_rank_loss, cur_rr = get_loss(prediction, gt_batch, price_batch, mask_batch,\n stock_num, alpha)\n loss += cur_loss.item()\n reg_loss += cur_reg_loss.item()\n rank_loss += cur_rank_loss.item()\n cur_valid_pred[:, cur_offset - (start_index - lookback_length - steps + 1)] = cur_rr[:, 0].cpu()\n cur_valid_gt[:, cur_offset - (start_index - lookback_length - steps + 1)] = gt_batch[:, 0].cpu()\n cur_valid_mask[:, cur_offset - (start_index - lookback_length - steps + 1)] = mask_batch[:, 0].cpu()\n loss = loss / (end_index - start_index)\n reg_loss = reg_loss / (end_index - start_index)\n rank_loss = rank_loss / (end_index - start_index)\n cur_valid_perf = evaluate(cur_valid_pred, cur_valid_gt, cur_valid_mask)\n return loss, reg_loss, rank_loss, cur_valid_perf\n\n\ndef get_batch(offset=None):\n if offset is None:\n offset = random.randrange(0, valid_index)\n seq_len = lookback_length\n mask_batch = mask_data[:, offset: offset + seq_len + steps]\n mask_batch = np.min(mask_batch, axis=1)\n return (\n eod_data[:, offset:offset + seq_len, :],\n np.expand_dims(mask_batch, axis=1),\n np.expand_dims(price_data[:, offset + seq_len - 1], axis=1),\n np.expand_dims(gt_data[:, offset + seq_len + steps - 1], axis=1))\n\n\nfor epoch in range(epochs):\n print(\"epoch{}##########################################################\".format(epoch + 1))\n np.random.shuffle(batch_offsets)\n tra_loss = 0.0\n tra_reg_loss = 0.0\n tra_rank_loss = 0.0\n for j in range(valid_index - lookback_length - steps + 1):\n data_batch, mask_batch, price_batch, gt_batch = map(\n lambda x: torch.Tensor(x).to(device),\n get_batch(batch_offsets[j])\n )\n optimizer.zero_grad()\n prediction = model(data_batch)\n cur_loss, cur_reg_loss, cur_rank_loss, _ = get_loss(prediction, gt_batch, price_batch, mask_batch,\n stock_num, alpha)\n cur_loss = cur_loss\n cur_loss.backward()\n optimizer.step()\n\n tra_loss += cur_loss.item()\n tra_reg_loss += cur_reg_loss.item()\n tra_rank_loss += cur_rank_loss.item()\n tra_loss = tra_loss / (valid_index - lookback_length - steps + 1)\n tra_reg_loss = tra_reg_loss / (valid_index - lookback_length - steps + 1)\n tra_rank_loss = tra_rank_loss / (valid_index - lookback_length - steps + 1)\n print('Train : loss:{:.2e} = {:.2e} + alpha*{:.2e}'.format(tra_loss, tra_reg_loss, tra_rank_loss))\n\n val_loss, val_reg_loss, val_rank_loss, val_perf = validate(valid_index, test_index)\n print('Valid : loss:{:.2e} = {:.2e} + alpha*{:.2e}'.format(val_loss, val_reg_loss, val_rank_loss))\n\n test_loss, test_reg_loss, test_rank_loss, test_perf = validate(test_index, trade_dates)\n print('Test: loss:{:.2e} = {:.2e} + alpha*{:.2e}'.format(test_loss, test_reg_loss, test_rank_loss))\n\n if val_loss < best_valid_loss:\n best_valid_loss = val_loss\n best_valid_perf = val_perf\n best_test_perf = test_perf\n\n print('Valid performance:\\n', 'mse:{:.2e}, IC:{:.2e}, RIC:{:.2e}, prec@10:{:.2e}, SR:{:.2e}'.format(val_perf['mse'], val_perf['IC'],\n val_perf['RIC'], val_perf['prec_10'], val_perf['sharpe5']))\n print('Test performance:\\n', 'mse:{:.2e}, IC:{:.2e}, RIC:{:.2e}, prec@10:{:.2e}, SR:{:.2e}'.format(test_perf['mse'], test_perf['IC'],\n test_perf['RIC'], test_perf['prec_10'], test_perf['sharpe5']), '\\n\\n')\n" } ]