Misc improvements for evaluation

Misc improvements for evaluation

 - restructure results
 - refactor Variations, VariationChain

#!/usr/bin/env python3
from collections import OrderedDict
import os
import json
import yaml
import re
from pathlib import Path
import redo
redo.ifchange('comap.py'); import comap

project = redo.base
project_path = Path('projects') / project

with redo.ifchange(evaluation_config=project_path / 'evaluation-config.yaml') \
        as files:
    evaluation_config = yaml.load(files['evaluation_config'])

databases = evaluation_config['databases']
outcome_ids = evaluation_config['outcome_ids']
concept_filenames = {
    outcome_id: '{}.{}.concepts.json'.format(project, outcome_id)
    for outcome_id in outcome_ids
}
child_concept_filenames = {
    outcome_id: '{}.{}.child-concepts.json'.format(project, outcome_id)
    for outcome_id in outcome_ids
}
casedef_paths = {
    outcome_id: project_path / 'case-definitions' / (outcome_id + '.yaml')
    for outcome_id in outcome_ids
}

with redo.ifchange(coding_systems='config/coding_systems.yaml',
                   references=(project_path / 'reference.yaml').as_posix(),
                   concepts=concept_filenames,
                   child_concepts=child_concept_filenames,
                   casedefs=casedef_paths) as files:
    coding_systems = yaml.load(files['coding_systems'])
    references = yaml.load(files['references'])
    concepts_by_outcome = {}
    for outcome_id, f in files['concepts'].items():
        concepts_by_outcome[outcome_id] = json.load(f)
    child_concepts_by_outcome = {}
    for outcome_id, f in files['child_concepts'].items():
        child_concepts_by_outcome[outcome_id] = json.load(f).items()
    outcomes = {
        outcome_id: yaml.load(f)
        for outcome_id, f in files['casedefs'].items()
    }

class AbstractVariation(object):

    def description(self):
        return self.__doc__

    def set_context(self, outcome_id, coding_system):
        self.outcome_id = outcome_id
        self.coding_system = coding_system

    def vary_codes(self, generated, reference):
        return generated, reference

    def vary_concepts(self, concepts, reference):
        return concepts

class VariationChain(AbstractVariation):

    """Chain of variations that are applied to the input left-to-right."""

    def __init__(self, *variations):
        self.variations = variations

    def set_context(self, *args, **kwargs):
        for variation in self.variations:
            variation.set_context(*args, **kwargs)

    def description(self):
        res = '; '.join(v.description() for v in self.variations)
        if res:
            return res[0].upper() + res[1:]
        else:
            return ""

    def vary_codes(self, generated, reference):
        for variation in self.variations:
            generated, reference = variation.vary_codes(generated, reference)
        return generated, reference

    def vary_concepts(self, concepts, reference):
        for variation in self.variations:
            concepts = variation.vary_concepts(concepts, reference)
        return concepts

class Identity(AbstractVariation):

    """Identity"""

    pass

class CodeNormalization(AbstractVariation):

    def __init__(self, coding_system_prefixes=['']):
        """`coding_systems` is None or a list of prefixes and this
        normalization is applied only if the list contains a prefix of the
        current coding systems."""
        self.coding_system_prefixes = coding_system_prefixes

    def vary_codes(self, generated, reference):
        if any(self.coding_system.startswith(prefix)
               for prefix in self.coding_system_prefixes):
            return [self.normalize(c) for c in generated], \
                [self.normalize(c) for c in reference]
        else:
            return generated, reference

class NormalizeCodeCase(CodeNormalization):

    """Normalize code case"""

    def normalize(self, code):
        return code.upper()

class NormalizeCodeSuffixDot(CodeNormalization):

    """Drop suffix dots in codes in ICPC"""

    def __init__(self):
        super().__init__(['ICPC'])

    def normalize(self, code):
        if code[-1] == '.':
            return code[:-1]
        else:
            return code

class NormalizeCodeXDD(CodeNormalization):

    """Normalize codes to XDD in ICD and ICPC"""

    def __init__(self):
        super().__init__(['ICD', 'ICPC'])

    def normalize(self, code):
        code = code.upper()
        m1 = re.match(comap.RegexHierarchy.TYPES['LDD'], code)
        try:
            return m1.group('parent')
        except:
            m2 = re.match(comap.RegexHierarchy.TYPES['DDD'], code)
            try:
                return m2.group('parent')
            except:
                return code

class NormalizeRead2Codes(CodeNormalization):

    """Normalize READ2 codes to 5-byte"""

    read_re = re.compile(r'(?P<code>[A-Z\d][A-Za-z0-9.]{4}).*')

    def __init__(self):
        super().__init__(['RCD2'])

    def normalize(self, code):
        m = self.read_re.match(code)
        assert m, "Not a READ2 code: {}".format(code)
        return m.group('code')

class IncludeRelatedCodes(AbstractVariation):

    """Include related codes"""

    hierarchies = {
        'ICD9CM': comap.RegexHierarchy('DDD'),
        'ICD10CM': comap.RegexHierarchy('LDD'),
        'ICPC': comap.RegexHierarchy('LDD'),
        'RCD2': comap.Read2Hierarchy(),
    }

    def __init__(self, direction):
        """ direction : 'FN-FP' | 'FN' """
        self.direction = direction

    def description(self):
        return super().description() + ' ({})'.format(self.direction)

    def vary_codes(self, generated, reference):
        hierarchy = self.hierarchies[self.coding_system]
        codes = comap.confusion_matrix(generated, reference)
        related_codes, derivated_codes = {
            'FN': comap.include_related,
            'FN-FP': comap.include_related_bidirect,
        }[self.direction](hierarchy, codes=codes)
        return comap.from_confusion_matrix(related_codes)

class IncludeChildConcepts(AbstractVariation):

    """Include child concepts"""

    def vary_concepts(self, concepts, reference):
        cuis = [c['cui'] for c in concepts]
        child_concepts_by_parent = child_concepts_by_outcome[self.outcome_id]
        child_concepts = [
            child_concept
            for parent_id, concepts_for_parent in child_concepts_by_parent
            for child_concept in concepts_for_parent
            if child_concept not in cuis
        ]
        return concepts + child_concepts

def create_results(references, concepts_by_outcome, variations):

    res = OrderedDict()
    for outcome_id in outcome_ids:
        res[outcome_id] = OrderedDict()
        for database, coding_system in databases:
            mapping = references[database]['mappings'][outcome_id]
            res[outcome_id][database] = OrderedDict([])

            if 'inclusion' in mapping:
                reference = mapping['inclusion'] + mapping.get('exclusion', [])
            else:
                print("No mapping for {} in {}".format(outcome_id, database))
                reference = None

            if reference is not None:
                res[outcome_id][database]['variations'] = list()
                for variation in variations:
                    variation.set_context(outcome_id, coding_system)
                    generated_concepts = variation.vary_concepts(
                        concepts_by_outcome[outcome_id], reference)
                    generated = [
                        code['id']
                        for concept in generated_concepts
                        for code in concept['sourceConcepts']
                        if code['codingSystem'] == coding_system
                    ]
                    varied_generated, varied_reference = \
                        variation.vary_codes(generated, reference)
                    codes = comap.confusion_matrix(generated=varied_generated,
                                                  reference=varied_reference)
                    measures = comap.measures(codes=codes)
                    for key, values in codes.items():
                        try:
                            sorted(values)
                        except:
                            print([v for v in values if type(v) != str])
                            raise
                    comparison = OrderedDict([
                        ('codes', OrderedDict([
                            (key, sorted(value))
                            for key, value in codes.items()
                        ])),
                        ('measures', OrderedDict([
                            (key, value) # and round(value, 2))
                            for key, value in measures.items()
                        ])),
                    ])
                    res[outcome_id][database]['variations'].append(OrderedDict([
                        ('name', variation.description()),
                        ('comparison', comparison),
                    ]))

            if 'comment' in mapping:
                res[outcome_id][database]['comment'] = mapping['comment']
    return res

default_variations = VariationChain(NormalizeCodeCase(), NormalizeCodeSuffixDot(), NormalizeRead2Codes())
variations = [
    Identity(),
    default_variations,
    VariationChain(default_variations, NormalizeCodeXDD()),
    VariationChain(default_variations, IncludeRelatedCodes('FN')),
    VariationChain(default_variations, IncludeRelatedCodes('FN-FP')),
    VariationChain(default_variations, IncludeChildConcepts()),
]

evaluation = create_results(references, concepts_by_outcome, variations)

# Output results as YAML

with redo.output() as f:
    yaml.dump(evaluation, f)

if os.getenv('DOCTEST'):
    import doctest
    doctest.testmod()
    doctest.testmod(comap)

#!/usr/bin/env python3
from collections import OrderedDict
import os
import json
import yaml
import re
from pathlib import Path
import redo
redo.ifchange('comap.py'); import comap

project = redo.base
project_path = Path('projects') / project

with redo.ifchange(evaluation_config=project_path / 'evaluation-config.yaml') \
        as files:
    evaluation_config = yaml.load(files['evaluation_config'])

databases = evaluation_config['databases']
outcome_ids = evaluation_config['outcome_ids']
concept_filenames = {
    outcome_id: '{}.{}.concepts.json'.format(project, outcome_id)
    for outcome_id in outcome_ids
}
child_concept_filenames = {
    outcome_id: '{}.{}.child-concepts.json'.format(project, outcome_id)
    for outcome_id in outcome_ids
}
casedef_paths = {
    outcome_id: project_path / 'case-definitions' / (outcome_id + '.yaml')
    for outcome_id in outcome_ids
}

with redo.ifchange(coding_systems='config/coding_systems.yaml',
                   references=(project_path / 'reference.yaml').as_posix(),
                   concepts=concept_filenames,
                   child_concepts=child_concept_filenames,
                   casedefs=casedef_paths) as files:
    coding_systems = yaml.load(files['coding_systems'])
    references = yaml.load(files['references'])
    concepts_by_outcome = {}
    for outcome_id, f in files['concepts'].items():
        concepts_by_outcome[outcome_id] = json.load(f)
    child_concepts_by_outcome = {}
    for outcome_id, f in files['child_concepts'].items():
        child_concepts_by_outcome[outcome_id] = json.load(f).items()
    outcomes = {
        outcome_id: yaml.load(f)
        for outcome_id, f in files['casedefs'].items()
    }

def create_results(references, concepts_by_outcome, variations):

    res = OrderedDict()
    for outcome_id in outcome_ids:
        res[outcome_id] = OrderedDict()
        for database, coding_system in databases:
            res[outcome_id][database] = list()

            mapping = references[database]['mappings'].get(outcome_id)
            if mapping is None:
                print("No mapping for {} in {}".format(outcome_id, database))
                reference = None
            else:
                reference = mapping['inclusion'] + mapping.get('exclusion', [])

            for make_variation in variations:
                if reference is not None:
                    variation = make_variation(outcome_id, coding_system)
                    generated_concepts = variation.vary_concepts(
                        concepts_by_outcome[outcome_id], reference)
                    generated = [
                        code['id']
                        for concept in generated_concepts
                        for code in concept['sourceConcepts']
                        if code['codingSystem'] == coding_system
                    ]
                    varied_generated, varied_reference = \
                        variation.vary_codes(generated, reference)
                    codes = comap.confusion_matrix(generated=varied_generated,
                                                  reference=varied_reference)
                    measures = comap.measures(codes=codes)
                    comparison = OrderedDict([
                        ('codes', OrderedDict([
                            (key, sorted(value))
                            for key, value in codes.items()
                        ])),
                        ('measures', OrderedDict([
                            (key, value and round(value, 2))
                            for key, value in measures.items()
                        ])),
                    ])
                else:
                    comparison = None
                res[outcome_id][database].append(OrderedDict([
                    ('name', variation.description()),
                    ('comparison', comparison),
                ]))
    return res

class AbstractVariation(object):

    def __init__(self, outcome_id, coding_system):
        self.outcome_id = outcome_id
        self.coding_system = coding_system

    def description(self):
        return self.__doc__

    def vary_codes(self, generated, reference):
        return generated, reference

    def vary_concepts(self, concepts, reference):
        return concepts

def variation_chain(*variations):

    class VariationChain(AbstractVariation):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.variations = [v(*args, **kwargs) for v in variations]

        def description(self):
            res = ' -> '.join(v.description().lower() for v in self.variations)
            if res:
                return res[0].upper() + res[1:]
            else:
                return ""

        def vary_codes(self, generated, reference):
            for variation in self.variations:
                generated, reference = variation.vary_codes(generated, reference)
            return generated, reference

        def vary_concepts(self, concepts, reference):
            for variation in self.variations:
                concepts = variation.vary_concepts(concepts, reference)
            return concepts

    return VariationChain

class Identity(AbstractVariation):

    """Identity"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

class CodeNormalization(AbstractVariation):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def vary_codes(self, generated, reference):
        return [self.normalize(c) for c in generated], \
            [self.normalize(c) for c in reference]

class NormalizeCodeCase(CodeNormalization):

    """Normalize code case"""

    @classmethod
    def normalize(cls, code):
        return code.upper()

class NormalizeCodeSuffixDot(CodeNormalization):

    """Drop suffix dots in codes"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def normalize(cls, code):
        if code[-1] == '.':
            return code[:-1]
        else:
            return code

class NormalizeCodeXDD(CodeNormalization):

    """Normalize codes (case and to XDD)"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @classmethod
    def normalize(cls, code):
        code = code.upper()
        m1 = re.match(comap.RegexHierarchy.TYPES['LDD'], code)
        try:
            return m1.group('parent')
        except:
            m2 = re.match(comap.RegexHierarchy.TYPES['DDD'], code)
            try:
                return m2.group('parent')
            except:
                return code

class IncludeRelatedCodes(AbstractVariation):

    """Include related codes"""

    hierarchies = {
        'ICD9CM': comap.RegexHierarchy('DDD'),
        'ICD10CM': comap.RegexHierarchy('LDD'),
        'ICPC': comap.RegexHierarchy('LDD'),
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def vary_codes(self, generated, reference):
        hierarchy = self.hierarchies[self.coding_system]
        codes = comap.confusion_matrix(generated, reference)
        related_codes, derivated_codes = comap.include_related(hierarchy, codes=codes)
        return comap.from_confusion_matrix(related_codes)

class IncludeRelatedCodesBidirect(AbstractVariation):

    """Include bidirectionally related codes"""

    hierarchies = {
        'ICD9CM': comap.RegexHierarchy('DDD'),
        'ICD10CM': comap.RegexHierarchy('LDD'),
        'ICPC': comap.RegexHierarchy('LDD'),
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def vary_codes(self, generated, reference):
        hierarchy = self.hierarchies[self.coding_system]
        codes = comap.confusion_matrix(generated, reference)
        related_codes, derivated_codes = comap.include_related_bidirect(hierarchy, codes=codes)
        return comap.from_confusion_matrix(related_codes)

class IncludeChildConcepts(AbstractVariation):

    """Include child concepts"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def vary_concepts(self, concepts, reference):
        cuis = [c['cui'] for c in concepts]
        child_concepts_by_parent = child_concepts_by_outcome[self.outcome_id]
        child_concepts = [
            child_concept
            for parent_id, concepts_for_parent in child_concepts_by_parent
            for child_concept in concepts_for_parent
            if child_concept not in cuis
        ]
        return concepts + child_concepts

variations = [
    Identity,
    variation_chain(NormalizeCodeCase, NormalizeCodeSuffixDot),
    variation_chain(NormalizeCodeCase, NormalizeCodeSuffixDot, NormalizeCodeXDD),
    variation_chain(NormalizeCodeCase, NormalizeCodeSuffixDot, IncludeRelatedCodes),
    variation_chain(NormalizeCodeCase, NormalizeCodeSuffixDot, IncludeRelatedCodesBidirect),
    variation_chain(NormalizeCodeCase, NormalizeCodeSuffixDot, IncludeChildConcepts),
]

evaluation = create_results(references, concepts_by_outcome, variations)

# Output results as YAML

with redo.output() as f:
    yaml.dump(evaluation, f)

if os.getenv('DOCTEST'):
    import doctest
    doctest.testmod()
    doctest.testmod(comap)


MT data and fix errors

MT data and fix errors

program make_basketball

make_hof

make_allstar

make_match

make_master

end

program read_mt

insheet using ${basketball}mt_bk_output.csv, clear

gen wikihandle = subinstr(answer, "http://en.wikipedia.org/wiki/","",.)
keep inputp inputdisp wikihandle

gen tmp = subinstr(inputd, " ","_",.)
fix_errors

save ${basketball}mt_bk_output, replace

end

program fix_errors

replace wiki = tmp if inputd == "Slater Martin"
replace wiki = tmp if inputd == "Tom Meschery"
replace wiki = tmp if wiki == "Micheal_Williams"
replace wiki = tmp if inputd == "Michael Cage"
replace wiki = tmp if inputd == "Alvin Robertson"
replace wiki = tmp if inputd == "Magic Johnson"
replace wiki = tmp if inputd == "Michael Jordan"
replace wiki = tmp if inputd == "Phil Jordon"

replace wiki = "George_L._Johnson" if inputd == "George L. Johnson"

replace wiki = "Wayne_Cooper_(basketball)" if inputd == "Wayne Cooper"

replace wiki = "" if inputd == "George L. Johnson"

replace wiki = "Micheal_Ray_Richardson" if inputd == "Micheal Ray Richardson"

replace wiki = "Eddie_Johnson_(basketball,_born_1955)" if inputd == "Eddie Johnson"
replace wiki = "Eddie_A._Johnson" if inputd == "Eddie A. Johnson"

replace wiki = "Mike_Brown_(basketball,_born_1963)" if inputd == "Mike Brown"

replace wiki = "Mahmoud_Abdul-Rauf" if wiki == "Mahmud"
replace wiki = "Jan_van_Breda_Kolff" if inputd == "Jan Vanbredakolff"
replace wiki = "Norm_Van_Lier" if inputd == "Norm Vanlier"
replace wiki = "Danny_Vranes" if inputd == "Danny Vranes"
replace wiki = "Dick_Barnett" if inputd == "Dick Barnett"
replace wiki = "Walt_Wesley" if inputd == "Walt Wesley"

end

// This makes baseball data for MT
program make_mt

use ${basketball}bk_master, clear

gen displayname = namefirst + " " + namelast
gen searchname = namefirst + "+" + namelast + "+" + " basketball"

gen link = "http://en.wikipedia.org/w/index.php?title=Special%3ASearch&profile=default&search=" + searchname + "&fulltext=Search"

outsheet playerid displayname link using ${stash}mt_input_bk.csv, replace comma

end

// This makes master baseball dataset
program make_master

insheet using ${basketball}players.csv, clear

replace college = college + ", " + birthdate if v12 != ""
replace birthdate = v12 if v12 != ""
drop v12

merge 1:1 ilkid using ${stash}bk_tmp, keep(match) nogen

gen birthyear = substr(birthdate,1,4)
rename firstseason debutyear
rename lastseason lastyear

rename ilkid playerid

keep playerid firstname lastname debutyear lastyear birthyear minutesrank minutes gp asts pts reb

rename firstname namefirst
rename lastname namelast

gen playername = namefirst + " " + namelast

// generate file for amazon to get me data on
drop if debutyear < 1940
drop if debutyear > 1990

drop if minutesrank > 1000

save ${basketball}bk_master, replace

end

program make_careermins

insheet using ${basketball}player_career.csv, clear
keep ilk firstname lastname gp minutes* pts asts reb turnover

// this data is at player-league level.
local vars "gp minutes pts reb asts turnover"
foreach x in `vars'{
    bysort ilk: egen tmp = total(`x')
    replace `x' = tmp
    drop tmp
}

bysort ilk: drop if _n>1

gsort- minutes
gen minutesrank = _n

save ${stash}bk_tmp, replace

end

program make_basketball

make_hof

make_allstar

make_match

make_master

end

// This makes baseball data for MT
program make_mt

use ${basketball}bk_master, clear

gen displayname = namefirst + " " + namelast
gen searchname = namefirst + "+" + namelast + "+" + " basketball"

gen link = "http://en.wikipedia.org/w/index.php?title=Special%3ASearch&profile=default&search=" + searchname + "&fulltext=Search"

outsheet playerid displayname link using ${stash}mt_input_bk.csv, replace comma

end

// This makes master baseball dataset
program make_master

insheet using ${basketball}players.csv, clear

replace college = college + ", " + birthdate if v12 != ""
replace birthdate = v12 if v12 != ""
drop v12

merge 1:1 ilkid using ${stash}bk_tmp, keep(match) nogen

gen birthyear = substr(birthdate,1,4)
rename firstseason debutyear
rename lastseason lastyear

rename ilkid playerid

keep playerid firstname lastname debutyear lastyear birthyear minutesrank minutes gp asts pts reb

rename firstname namefirst
rename lastname namelast

gen playername = namefirst + " " + namelast

// generate file for amazon to get me data on
drop if debutyear < 1940
drop if debutyear > 1990

drop if minutesrank > 1000

save ${basketball}bk_master, replace

end

program make_careermins

insheet using ${basketball}player_career.csv, clear
keep ilk firstname lastname gp minutes* pts asts reb turnover

// this data is at player-league level.
local vars "gp minutes pts reb asts turnover"
foreach x in `vars'{
    bysort ilk: egen tmp = total(`x')
    replace `x' = tmp
    drop tmp
}

bysort ilk: drop if _n>1

gsort- minutes
gen minutesrank = _n

save ${stash}bk_tmp, replace

end


* changed FTP directory

* changed FTP directory

redo-ifchange $1

exec >&2  # redirect stdout to stderr
ftp -vi dieweltistgarnichtso.net << EOF
cd hosts/warumnicht.dieweltistgarnichtso.net
hash
put $1
bye

redo-ifchange $1

exec >&2  # redirect stdout to stderr
ftp -vi dieweltistgarnichtso.net << EOF
cd hosts/warumnicht.so
hash
put $1
bye


(minor) Tweak debug messages

(minor) Tweak debug messages

*! version 5.2.1 18jul2018

program reghdfe, eclass
	* Intercept old+version
	cap syntax, version old
	if !c(rc) {
		reghdfe_old, version
		exit
	}

	* Intercept old
	cap syntax anything(everything) [fw aw pw/], [*] old
	if !c(rc) {
		di as error "(running historical version of reghdfe)"
		if ("`weight'"!="") local weightexp [`weight'=`exp']
		reghdfe_old `anything' `weightexp', `options'
		exit
	}

	* Aux. subcommands
	cap syntax, [*]
	if inlist("`options'", "check", "compile", "reload", "update", "version", "requirements") {
		if ("`options'"=="compile") loc args force
		if ("`options'"=="check") loc options compile
		if ("`options'"=="update") {
			loc args 1
			loc options reload
		}
		loc subcmd = proper("`options'")
		`subcmd' `args'
	}
	else if replay() {
		Replay `0'
	}
	else {
		Cleanup 0
		Compile // takes 0.01s to run this useful check (ensures .mlib exists)
		cap noi Estimate `0'
		Cleanup `c(rc)'
	}
end

program Compile
	args force

	* Check dependencies
	ftools, check // in case lftools.mlib does not exist or is outdated
	ms_get_version ftools, min_version("2.30.1")
	ms_get_version reghdfe // save local package_version
	loc list_objects "FixedEffects() fixed_effects() BipartiteGraph()"
	loc list_functions "reghdfe_*() transform_*() accelerate_*() panelmean() panelsolve_*() lsmr()"
	loc list_misc "weighted_quadcolsum() safe_divide() check_convergence() precompute_inv_xx() _st_data_wrapper()"
	// TODO: prefix everything with reghdfe_*

	ms_compile_mata, ///
		package(reghdfe) ///
		version(`package_version') ///
		fun("`list_objects' `list_functions' `list_misc'") ///
		verbose ///
		`force'
end

program Reload
	* Internal debugging tool.
	* Updates dependencies and reghdfe from local path or from github
	* Usage:
	* 	reghdfe, update // from c:\git\..
	* 	reghdfe, reload // from github

	args online
	if ("`online'" == "") loc online 0

	di as text _n "{bf:reghdfe: updating required packages}"
	di as text "{hline 64}"

	* -ftools- https://github.com/sergiocorreia/ftools/
	cap ado uninstall ftools
	if (`online') net install ftools, from("https://github.com/sergiocorreia/ftools/raw/master/src/")
	if (!`online') net install ftools, from("c:\git\ftools\src")
	di as text "{hline 64}"
	ftools, compile
	di as text "{hline 64}"

	* Update -reghdfe-
	di as text _n  _n "{bf:reghdfe: updating self}"
	di as text "{hline 64}"
	qui ado uninstall reghdfe
	if (`online') net install reghdfe, from("https://github.com/sergiocorreia/reghdfe/raw/master/src/")
	if (!`online') net install reghdfe, from("c:\git\reghdfe\src")
	qui which reghdfe
	di as text "{hline 64}"
	reghdfe, compile
	di as text "{hline 64}"

	* Cleaning up
	di as text _n "{bf:Note:} You need to run {stata program drop _all} now."
end

program Version
	which reghdfe
	Requirements
end

program Requirements
	di as text _n "Required packages installed?"
	loc reqs ftools
	// ivreg2 avar tuples group3hdfe
	if (c(stata_version)<13) loc reqs `reqs' boottest

	loc ftools_github "https://github.com/sergiocorreia/ftools/raw/master/src/"

	loc error 0

	foreach req of local reqs {
		loc fn `req'.ado
		cap findfile `fn'
		if (_rc) {
			loc error 1
			di as text "{lalign 20:- `req'}" as error "not" _c
			di as text "    {stata ssc install `req':install from SSC}" _c
			if inlist("`req'", "ftools") {
				loc github ``req'_github'
				di as text `"    {stata `"net install `req', from(`"`github'"')"':install from github}"'
			}
			else {
				di as text // newline
			}
		}
		else {
			di as text "{lalign 20:- `req'}" as text "yes"
		}
	}

	if (`error') exit 601
end

program Cleanup
	args rc
	cap mata: mata drop HDFE
	cap mata: mata drop hdfe_*
	cap drop __temp_reghdfe_resid__
	cap matrix drop reghdfe_statsmatrix
	if (`rc') exit `rc'
end

program Parse
	* Trim whitespace (caused by "///" line continuations; aesthetic only)
	mata: st_local("0", stritrim(st_local("0")))

	* Main syntax
	#d;
	syntax varlist(fv ts numeric) [if] [in] [aw pw fw/] , [

		/* Model */
		Absorb(string) NOAbsorb
		SUmmarize SUmmarize2(string asis) /* simulate implicit options */

		/* Standard Errors */
		VCE(string) CLuster(string)

		/* Diagnostic */
		Verbose(numlist min=1 max=1 >=-1 <=5 integer)
		TIMEit

		/* Speedup and memory Tricks */
		NOSAMPle /* do not save e(sample) */
		COMPACT /* use as little memory as possible but is slower */

		/* Undocumented */
		KEEPSINgletons
		OLD /* use latest v3 */
		NOTES(string) /* NOTES(key=value ...), will be stored on e() */

		] [*] /* capture optimization options, display options, etc. */
		;
	#d cr

	* Unused
	* SAVEcache
	* USEcache
	* CLEARcache

	* Convert options to boolean
	if ("`verbose'" == "") loc verbose 0
	loc timeit = ("`timeit'"!="")
	loc drop_singletons = ("`keepsingletons'" == "")
	loc compact = ("`compact'" != "")

	if (`timeit') timer on 29

	* Sanity checks
	if (`verbose'>-1 & "`keepsingletons'"!="") {
		loc url "http://scorreia.com/reghdfe/nested_within_cluster.pdf"
		loc msg "WARNING: Singleton observations not dropped; statistical significance is biased"
		di as error `"`msg' {browse "`url'":(link)}"'
	}
	if ("`cluster'"!="") {
		_assert ("`vce'"==""), msg("cannot specify both cluster() and vce()")
		loc vce cluster `cluster'
		loc cluster // clear it to avoid bugs in subsequent lines
	}

	* Split varlist into <depvar> and <indepvars>
	ms_parse_varlist `varlist'
	if (`verbose' > 0) {
		di as text _n "## Parsing varlist: {res}`varlist'"
		return list
	}
	loc depvar `r(depvar)'
	loc indepvars `r(indepvars)'
	loc fe_format "`r(fe_format)'"
	loc basevars `r(basevars)'

	* Parse Weights
	if ("`weight'"!="") {
		unab exp : `exp', min(1) max(1) // simple weights only
	}

	* Parse VCE
	ms_parse_vce, vce(`vce') weighttype(`weight')
	if (`verbose' > 0) {
		di as text _n "## Parsing vce({res}`vce'{txt})"
		sreturn list
	}
	loc vcetype = "`s(vcetype)'"
	loc num_clusters = `s(num_clusters)'
	loc clustervars = "`s(clustervars)'"
	loc base_clustervars = "`s(base_clustervars)'"
	loc vceextra = "`s(vceextra)'"

	* Select sample (except for absvars)
	loc varlist `depvar' `indepvars' `base_clustervars'
	tempvar touse
	marksample touse, strok // based on varlist + cluster + if + in + weight

	* Parse noabsorb
	_assert  ("`absorb'`noabsorb'" != ""), msg("option {bf:absorb()} or {bf:noabsorb} required")
	if ("`noabsorb'" != "") {
		_assert ("`absorb'" == ""), msg("{bf:absorb()} and {bf:noabsorb} are mutually exclusive")
		tempvar c
		gen byte `c' = 1
		loc absorb `c'
	}

	if (`timeit') timer off 29

	* Construct HDFE object
	// SYNTAX: fixed_effects(absvars | , touse, wtype, wtvar, dropsing, verbose)
	mata: st_local("comma", strpos(`"`absorb'"', ",") ? "" : ",")
	if (`timeit') timer on 20
	mata: HDFE = fixed_effects(`"`absorb' `comma' `options'"', "`touse'", "`weight'", "`exp'", `drop_singletons', `verbose')
	if (`timeit') timer off 20
	mata: HDFE.cmdline = "reghdfe " + st_local("0")
	loc options `s(options)'

	mata: st_local("N", strofreal(HDFE.N))
	if (`N' == 0) error 2000

	* Fill out HDFE object
	* mata: HDFE.varlist = "`base_varlist'"
	mata: HDFE.depvar = "`depvar'"
	mata: HDFE.indepvars = "`indepvars'"
	mata: HDFE.vcetype = "`vcetype'"
	mata: HDFE.num_clusters = `num_clusters'
	mata: HDFE.clustervars = tokens("`clustervars'")
	mata: HDFE.base_clustervars = tokens("`base_clustervars'")
	mata: HDFE.vceextra = "`vceextra'"

	* Preserve memory
	mata: HDFE.compact = `compact'
	if (`compact') {
		loc panelvar "`_dta[_TSpanel]'"
		loc timevar "`_dta[_TStvar]'"
		mata: HDFE.panelvar = "`panelvar'"
		mata: HDFE.timevar = "`timevar'"
		c_local keepvars `basevars' `base_clustervars' `panelvar' `timevar' // `exp'
	}

	* Parse summarize
	if ("`summarize'" != "") {
		_assert ("`summarize2'" == ""), msg("summarize() syntax error")
		loc summarize2 mean min max  // default values
	}
	ParseSummarize `summarize2'
	mata: HDFE.summarize_stats = "`s(stats)'"
	mata: HDFE.summarize_quietly = `s(quietly)'

	* Parse misc options
	mata: HDFE.notes = `"`notes'"'
	mata: HDFE.store_sample = ("`nosample'"=="")
	mata: HDFE.timeit = `timeit'

	* Parse Coef Table Options (do this last!)
	_get_diopts diopts options, `options' // store in `diopts', and the rest back to `options'
	_assert (`"`options'"'==""), msg(`"invalid options: `options'"')
	if ("`hascons'"!="") di in ye "(option ignored: `hascons')"
	if ("`tsscons'"!="") di in ye "(option ignored: `tsscons')"
	mata: HDFE.diopts = `"`diopts'"'
end

program ParseSummarize, sclass
	sreturn clear
	syntax [namelist(name=stats)] , [QUIetly]
	local quietly = ("`quietly'"!="")
	sreturn loc stats "`stats'"
	sreturn loc quietly = `quietly'
end

// --------------------------------------------------------------------------

program Estimate, eclass
	ereturn clear

	* Parse and fill out HDFE object
	Parse `0'
	mata: st_local("timeit", strofreal(HDFE.timeit))
	mata: st_local("compact", strofreal(HDFE.compact))
	mata: st_local("verbose", strofreal(HDFE.verbose))

	* Compute degrees-of-freedom
	if (`timeit') timer on 21
	mata: HDFE.estimate_dof()
	if (`timeit') timer off 21

	* Save updated e(sample) (singletons reduce sample);
	* required to parse factor variables to partial out
	if (`timeit') timer on 29
	tempvar touse
	mata: HDFE.save_touse("`touse'")
	if (`timeit') timer off 29

	* Expand varlists
	if (`timeit') timer on 22
	mata: st_local("depvar", HDFE.depvar)
	mata: st_local("indepvars", HDFE.indepvars)
	if (`verbose' > 0) di as text _n "## Parsing and expanding indepvars: {res}`indepvars'"
	ms_expand_varlist `indepvars' if `touse'
	if (`verbose' > 0) return list
	mata: HDFE.fullindepvars = "`r(fullvarlist)'"
	mata: HDFE.indepvars = "`r(varlist)'"
	mata: HDFE.not_basevar = strtoreal(tokens("`r(not_omitted)'"))
	mata: HDFE.varlist = "`depvar' `r(varlist)'"
	if (`timeit') timer off 22

	* Stats
	mata: st_local("stats", HDFE.summarize_stats)
	if ("`stats'" != "") Stats `touse'

	* Condition number
	mata: HDFE.estimate_cond()

	* Preserve
	if (`compact') {
		if (`verbose' > 0) di as text "## Preserving dataset"
		preserve
		keep `keepvars'
	}

	* Partial out; save TSS of depvar
	if (`timeit') timer on 23
	mata: hdfe_variables = HDFE.partial_out(HDFE.varlist, 1) // 1=Save TSS of first var if HDFE.tss is missing
	if (`timeit') timer off 23

	* Regress
	if (`timeit') timer on 24
	tempname b V N rank df_r
	mata: reghdfe_post_ols(HDFE, hdfe_variables, "`b'", "`V'", "`N'", "`rank'", "`df_r'")
	mata: hdfe_variables = .
	* Restore
	if (`compact') {
		if (`verbose' > 0) di as text "## Restoring dataset"
		restore
		mata: st_local("residuals", HDFE.residuals)
		if ("`residuals'" != "") mata: HDFE.save_variable(HDFE.residuals, HDFE.residuals_vector, "Residuals")
	}
	RegressOLS `touse' `b' `V' `N' `rank' `df_r'
	if (`timeit') timer off 24

	* (optional) Store FEs
	if (`timeit') timer on 29
	reghdfe_store_alphas
	if (`timeit') timer off 29

	* View estimation tables
	mata: st_local("diopts", HDFE.diopts)
	Replay, `diopts'

	if (`timeit') {
		di as text _n "{bf: Timer results:}"
		timer list
		di as text "Legend: 20: Create HDFE object; 21: Estimate DoF; 22: expand varlists; 23: partial out; 24: regress; 29: rest"
		di
	}
end

program RegressOLS, eclass
	args touse b V N rank df_r

	mata: st_local("store_sample", strofreal(HDFE.store_sample))
	if (`store_sample') loc esample "esample(`touse')"

	mata: st_local("indepvars", HDFE.fullindepvars)
	if ("`indepvars'" != "") {
		matrix colnames `b' = `indepvars'
		matrix colnames `V' = `indepvars'
		matrix rownames `V' = `indepvars'
		_ms_findomitted `b' `V'
		ereturn post `b' `V', `esample' buildfvinfo depname(`depvar')
	}
	else {
		ereturn post, `esample' buildfvinfo depname(`depvar')
	}

	ereturn scalar N       = `N'
	ereturn scalar rank    = `rank'
	ereturn scalar df_r    = `df_r'
	ereturn local  cmd     "reghdfe"
	mata: HDFE.post()

	* Post stats
	cap conf matrix reghdfe_statsmatrix
	if (!c(rc)) {
		ereturn matrix summarize = reghdfe_statsmatrix
		mata: st_local("summarize_quietly", strofreal(HDFE.summarize_quietly))
		ereturn scalar summarize_quietly = `summarize_quietly'
	}
end

program Replay, rclass
	syntax [, *]

	if `"`e(cmd)'"' != "reghdfe"  {
	        error 301
	}

	_get_diopts options, `options'
	reghdfe_header // _coef_table_header
	di ""
	_coef_table, `options' // ereturn display, `options'
	return add // adds r(level), r(table), etc. to ereturn (before the footnote deletes them)
	reghdfe_footnote

	* Replay stats
	if (e(summarize_quietly)==0) {
		di as text _n "{sf:Regression Summary Statistics:}" _c
		matlist e(summarize)', border(top bottom) rowtitle(Variable) // twidth(18)
	}
end

program Stats
	args touse
	* Optional weights
	mata: st_local("weight", sprintf("[%s=%s]", HDFE.weight_type, HDFE.weight_var))
	assert "`weight'" != ""
	if ("`weight'" == "[=]") loc weight
	loc weight : subinstr local weight "[pweight" "[aweight"

	mata: st_local("stats", HDFE.summarize_stats)
	mata: st_local("varlist", HDFE.varlist)
	mata: st_local("cvars", invtokens(HDFE.cvars))
	loc full_varlist `varlist' `cvars'

	* quick workaround b/c -tabstat- does not support factor variables
	fvrevar `full_varlist', list
	loc full_varlist `r(varlist)'

	qui tabstat `full_varlist' if `touse' `weight' , stat(`stats') col(stat) save
	matrix reghdfe_statsmatrix = r(StatTotal)
end

*! version 5.2.1 18jul2018

program reghdfe, eclass
	* Intercept old+version
	cap syntax, version old
	if !c(rc) {
		reghdfe_old, version
		exit
	}

	* Intercept old
	cap syntax anything(everything) [fw aw pw/], [*] old
	if !c(rc) {
		di as error "(running historical version of reghdfe)"
		if ("`weight'"!="") local weightexp [`weight'=`exp']
		reghdfe_old `anything' `weightexp', `options'
		exit
	}

	* Aux. subcommands
	cap syntax, [*]
	if inlist("`options'", "check", "compile", "reload", "update", "version", "requirements") {
		if ("`options'"=="compile") loc args force
		if ("`options'"=="check") loc options compile
		if ("`options'"=="update") {
			loc args 1
			loc options reload
		}
		loc subcmd = proper("`options'")
		`subcmd' `args'
	}
	else if replay() {
		Replay `0'
	}
	else {
		Cleanup 0
		Compile // takes 0.01s to run this useful check (ensures .mlib exists)
		cap noi Estimate `0'
		Cleanup `c(rc)'
	}
end

program Compile
	args force

	* Check dependencies
	ftools, check // in case lftools.mlib does not exist or is outdated
	ms_get_version ftools, min_version("2.30.1")
	ms_get_version reghdfe // save local package_version
	loc list_objects "FixedEffects() fixed_effects() BipartiteGraph()"
	loc list_functions "reghdfe_*() transform_*() accelerate_*() panelmean() panelsolve_*() lsmr()"
	loc list_misc "weighted_quadcolsum() safe_divide() check_convergence() precompute_inv_xx() _st_data_wrapper()"
	// TODO: prefix everything with reghdfe_*

	ms_compile_mata, ///
		package(reghdfe) ///
		version(`package_version') ///
		fun("`list_objects' `list_functions' `list_misc'") ///
		verbose ///
		`force'
end

program Reload
	* Internal debugging tool.
	* Updates dependencies and reghdfe from local path or from github
	* Usage:
	* 	reghdfe, update // from c:\git\..
	* 	reghdfe, reload // from github

	args online
	if ("`online'" == "") loc online 0

	di as text _n "{bf:reghdfe: updating required packages}"
	di as text "{hline 64}"

	* -ftools- https://github.com/sergiocorreia/ftools/
	cap ado uninstall ftools
	if (`online') net install ftools, from("https://github.com/sergiocorreia/ftools/raw/master/src/")
	if (!`online') net install ftools, from("c:\git\ftools\src")
	di as text "{hline 64}"
	ftools, compile
	di as text "{hline 64}"

	* Update -reghdfe-
	di as text _n  _n "{bf:reghdfe: updating self}"
	di as text "{hline 64}"
	qui ado uninstall reghdfe
	if (`online') net install reghdfe, from("https://github.com/sergiocorreia/reghdfe/raw/master/src/")
	if (!`online') net install reghdfe, from("c:\git\reghdfe\src")
	qui which reghdfe
	di as text "{hline 64}"
	reghdfe, compile
	di as text "{hline 64}"

	* Cleaning up
	di as text _n "{bf:Note:} You need to run {stata program drop _all} now."
end

program Version
	which reghdfe
	Requirements
end

program Requirements
	di as text _n "Required packages installed?"
	loc reqs ftools
	// ivreg2 avar tuples group3hdfe
	if (c(stata_version)<13) loc reqs `reqs' boottest

	loc ftools_github "https://github.com/sergiocorreia/ftools/raw/master/src/"

	loc error 0

	foreach req of local reqs {
		loc fn `req'.ado
		cap findfile `fn'
		if (_rc) {
			loc error 1
			di as text "{lalign 20:- `req'}" as error "not" _c
			di as text "    {stata ssc install `req':install from SSC}" _c
			if inlist("`req'", "ftools") {
				loc github ``req'_github'
				di as text `"    {stata `"net install `req', from(`"`github'"')"':install from github}"'
			}
			else {
				di as text // newline
			}
		}
		else {
			di as text "{lalign 20:- `req'}" as text "yes"
		}
	}

	if (`error') exit 601
end

program Cleanup
	args rc
	cap mata: mata drop HDFE
	cap mata: mata drop hdfe_*
	cap drop __temp_reghdfe_resid__
	cap matrix drop reghdfe_statsmatrix
	if (`rc') exit `rc'
end

program Parse
	* Trim whitespace (caused by "///" line continuations; aesthetic only)
	mata: st_local("0", stritrim(st_local("0")))

	* Main syntax
	#d;
	syntax varlist(fv ts numeric) [if] [in] [aw pw fw/] , [

		/* Model */
		Absorb(string) NOAbsorb
		SUmmarize SUmmarize2(string asis) /* simulate implicit options */

		/* Standard Errors */
		VCE(string) CLuster(string)

		/* Diagnostic */
		Verbose(numlist min=1 max=1 >=-1 <=5 integer)
		TIMEit

		/* Speedup and memory Tricks */
		NOSAMPle /* do not save e(sample) */
		COMPACT /* use as little memory as possible but is slower */

		/* Undocumented */
		KEEPSINgletons
		OLD /* use latest v3 */
		NOTES(string) /* NOTES(key=value ...), will be stored on e() */

		] [*] /* capture optimization options, display options, etc. */
		;
	#d cr

	* Unused
	* SAVEcache
	* USEcache
	* CLEARcache

	* Convert options to boolean
	if ("`verbose'" == "") loc verbose 0
	loc timeit = ("`timeit'"!="")
	loc drop_singletons = ("`keepsingletons'" == "")
	loc compact = ("`compact'" != "")

	if (`timeit') timer on 29

	* Sanity checks
	if (`verbose'>-1 & "`keepsingletons'"!="") {
		loc url "http://scorreia.com/reghdfe/nested_within_cluster.pdf"
		loc msg "WARNING: Singleton observations not dropped; statistical significance is biased"
		di as error `"`msg' {browse "`url'":(link)}"'
	}
	if ("`cluster'"!="") {
		_assert ("`vce'"==""), msg("cannot specify both cluster() and vce()")
		loc vce cluster `cluster'
		loc cluster // clear it to avoid bugs in subsequent lines
	}

	* Split varlist into <depvar> and <indepvars>
	ms_parse_varlist `varlist'
	if (`verbose' > 0) return list
	loc depvar `r(depvar)'
	loc indepvars `r(indepvars)'
	loc fe_format "`r(fe_format)'"
	loc basevars `r(basevars)'

	* Parse Weights
	if ("`weight'"!="") {
		unab exp : `exp', min(1) max(1) // simple weights only
	}

	* Parse VCE
	ms_parse_vce, vce(`vce') weighttype(`weight')
	loc vcetype = "`s(vcetype)'"
	loc num_clusters = `s(num_clusters)'
	loc clustervars = "`s(clustervars)'"
	loc base_clustervars = "`s(base_clustervars)'"
	loc vceextra = "`s(vceextra)'"

	* Select sample (except for absvars)
	loc varlist `depvar' `indepvars' `base_clustervars'
	tempvar touse
	marksample touse, strok // based on varlist + cluster + if + in + weight

	* Parse noabsorb
	_assert  ("`absorb'`noabsorb'" != ""), msg("option {bf:absorb()} or {bf:noabsorb} required")
	if ("`noabsorb'" != "") {
		_assert ("`absorb'" == ""), msg("{bf:absorb()} and {bf:noabsorb} are mutually exclusive")
		tempvar c
		gen byte `c' = 1
		loc absorb `c'
	}

	if (`timeit') timer off 29

	* Construct HDFE object
	// SYNTAX: fixed_effects(absvars | , touse, wtype, wtvar, dropsing, verbose)
	mata: st_local("comma", strpos(`"`absorb'"', ",") ? "" : ",")
	if (`timeit') timer on 20
	mata: HDFE = fixed_effects(`"`absorb' `comma' `options'"', "`touse'", "`weight'", "`exp'", `drop_singletons', `verbose')
	if (`timeit') timer off 20
	mata: HDFE.cmdline = "reghdfe " + st_local("0")
	loc options `s(options)'

	mata: st_local("N", strofreal(HDFE.N))
	if (`N' == 0) error 2000

	* Fill out HDFE object
	* mata: HDFE.varlist = "`base_varlist'"
	mata: HDFE.depvar = "`depvar'"
	mata: HDFE.indepvars = "`indepvars'"
	mata: HDFE.vcetype = "`vcetype'"
	mata: HDFE.num_clusters = `num_clusters'
	mata: HDFE.clustervars = tokens("`clustervars'")
	mata: HDFE.base_clustervars = tokens("`base_clustervars'")
	mata: HDFE.vceextra = "`vceextra'"

	* Preserve memory
	mata: HDFE.compact = `compact'
	if (`compact') {
		loc panelvar "`_dta[_TSpanel]'"
		loc timevar "`_dta[_TStvar]'"
		mata: HDFE.panelvar = "`panelvar'"
		mata: HDFE.timevar = "`timevar'"
		c_local keepvars `basevars' `base_clustervars' `panelvar' `timevar' // `exp'
	}

	* Parse summarize
	if ("`summarize'" != "") {
		_assert ("`summarize2'" == ""), msg("summarize() syntax error")
		loc summarize2 mean min max  // default values
	}
	ParseSummarize `summarize2'
	mata: HDFE.summarize_stats = "`s(stats)'"
	mata: HDFE.summarize_quietly = `s(quietly)'

	* Parse misc options
	mata: HDFE.notes = `"`notes'"'
	mata: HDFE.store_sample = ("`nosample'"=="")
	mata: HDFE.timeit = `timeit'

	* Parse Coef Table Options (do this last!)
	_get_diopts diopts options, `options' // store in `diopts', and the rest back to `options'
	_assert (`"`options'"'==""), msg(`"invalid options: `options'"')
	if ("`hascons'"!="") di in ye "(option ignored: `hascons')"
	if ("`tsscons'"!="") di in ye "(option ignored: `tsscons')"
	mata: HDFE.diopts = `"`diopts'"'
end

program ParseSummarize, sclass
	sreturn clear
	syntax [namelist(name=stats)] , [QUIetly]
	local quietly = ("`quietly'"!="")
	sreturn loc stats "`stats'"
	sreturn loc quietly = `quietly'
end

// --------------------------------------------------------------------------

program Estimate, eclass
	ereturn clear

	* Parse and fill out HDFE object
	Parse `0'
	mata: st_local("timeit", strofreal(HDFE.timeit))
	mata: st_local("compact", strofreal(HDFE.compact))
	mata: st_local("verbose", strofreal(HDFE.verbose))

	* Compute degrees-of-freedom
	if (`timeit') timer on 21
	mata: HDFE.estimate_dof()
	if (`timeit') timer off 21

	* Save updated e(sample) (singletons reduce sample);
	* required to parse factor variables to partial out
	if (`timeit') timer on 29
	tempvar touse
	mata: HDFE.save_touse("`touse'")
	if (`timeit') timer off 29

	* Expand varlists
	if (`timeit') timer on 22
	mata: st_local("depvar", HDFE.depvar)
	mata: st_local("indepvars", HDFE.indepvars)
	if (`verbose' > 0) di as text _n "## Parsing and expanding indepvars: {res}`indepvars'"
	ms_expand_varlist `indepvars' if `touse'
	if (`verbose' > 0) return list
	mata: HDFE.fullindepvars = "`r(fullvarlist)'"
	mata: HDFE.indepvars = "`r(varlist)'"
	mata: HDFE.not_basevar = strtoreal(tokens("`r(not_omitted)'"))
	mata: HDFE.varlist = "`depvar' `r(varlist)'"
	if (`timeit') timer off 22

	* Stats
	mata: st_local("stats", HDFE.summarize_stats)
	if ("`stats'" != "") Stats `touse'

	* Condition number
	mata: HDFE.estimate_cond()

	* Preserve
	if (`compact') {
		if (`verbose' > 0) di as text "## Preserving dataset"
		preserve
		keep `keepvars'
	}

	* Partial out; save TSS of depvar
	if (`timeit') timer on 23
	mata: hdfe_variables = HDFE.partial_out(HDFE.varlist, 1) // 1=Save TSS of first var if HDFE.tss is missing
	if (`timeit') timer off 23

	* Regress
	if (`timeit') timer on 24
	tempname b V N rank df_r
	mata: reghdfe_post_ols(HDFE, hdfe_variables, "`b'", "`V'", "`N'", "`rank'", "`df_r'")
	mata: hdfe_variables = .
	* Restore
	if (`compact') {
		if (`verbose' > 0) di as text "## Restoring dataset"
		restore
		mata: st_local("residuals", HDFE.residuals)
		if ("`residuals'" != "") mata: HDFE.save_variable(HDFE.residuals, HDFE.residuals_vector, "Residuals")
	}
	RegressOLS `touse' `b' `V' `N' `rank' `df_r'
	if (`timeit') timer off 24

	* (optional) Store FEs
	if (`timeit') timer on 29
	reghdfe_store_alphas
	if (`timeit') timer off 29

	* View estimation tables
	mata: st_local("diopts", HDFE.diopts)
	Replay, `diopts'

	if (`timeit') {
		di as text _n "{bf: Timer results:}"
		timer list
		di as text "Legend: 20: Create HDFE object; 21: Estimate DoF; 22: expand varlists; 23: partial out; 24: regress; 29: rest"
		di
	}
end

program RegressOLS, eclass
	args touse b V N rank df_r

	mata: st_local("store_sample", strofreal(HDFE.store_sample))
	if (`store_sample') loc esample "esample(`touse')"

	mata: st_local("indepvars", HDFE.fullindepvars)
	if ("`indepvars'" != "") {
		matrix colnames `b' = `indepvars'
		matrix colnames `V' = `indepvars'
		matrix rownames `V' = `indepvars'
		_ms_findomitted `b' `V'
		ereturn post `b' `V', `esample' buildfvinfo depname(`depvar')
	}
	else {
		ereturn post, `esample' buildfvinfo depname(`depvar')
	}

	ereturn scalar N       = `N'
	ereturn scalar rank    = `rank'
	ereturn scalar df_r    = `df_r'
	ereturn local  cmd     "reghdfe"
	mata: HDFE.post()

	* Post stats
	cap conf matrix reghdfe_statsmatrix
	if (!c(rc)) {
		ereturn matrix summarize = reghdfe_statsmatrix
		mata: st_local("summarize_quietly", strofreal(HDFE.summarize_quietly))
		ereturn scalar summarize_quietly = `summarize_quietly'
	}
end

program Replay, rclass
	syntax [, *]

	if `"`e(cmd)'"' != "reghdfe"  {
	        error 301
	}

	_get_diopts options, `options'
	reghdfe_header // _coef_table_header
	di ""
	_coef_table, `options' // ereturn display, `options'
	return add // adds r(level), r(table), etc. to ereturn (before the footnote deletes them)
	reghdfe_footnote

	* Replay stats
	if (e(summarize_quietly)==0) {
		di as text _n "{sf:Regression Summary Statistics:}" _c
		matlist e(summarize)', border(top bottom) rowtitle(Variable) // twidth(18)
	}
end

program Stats
	args touse
	* Optional weights
	mata: st_local("weight", sprintf("[%s=%s]", HDFE.weight_type, HDFE.weight_var))
	assert "`weight'" != ""
	if ("`weight'" == "[=]") loc weight
	loc weight : subinstr local weight "[pweight" "[aweight"

	mata: st_local("stats", HDFE.summarize_stats)
	mata: st_local("varlist", HDFE.varlist)
	mata: st_local("cvars", invtokens(HDFE.cvars))
	loc full_varlist `varlist' `cvars'

	* quick workaround b/c -tabstat- does not support factor variables
	fvrevar `full_varlist', list
	loc full_varlist `r(varlist)'

	qui tabstat `full_varlist' if `touse' `weight' , stat(`stats') col(stat) save
	matrix reghdfe_statsmatrix = r(StatTotal)
end


ENH: Update mat2nparray with fixes from gist

capture program drop mat2nparray
program define mat2nparray
    version 11.0
    syntax namelist(min=1), SAVing(str) [ Format(str) APPend REPlace ]
    if "`format'"=="" local format "%16.0g"
    local saving: subinstr local saving "." ".", count(local ext)
    if !`ext' local saving "`saving'.py"
    tempname myfile
    file open `myfile' using "`saving'", write text `append' `replace'
    file write `myfile' "import numpy as np" _n _n

    foreach mat of local namelist {
        mkarray `mat' `myfile' `format'
    }
    file write `myfile' "class Bunch(dict):" _n
    file write `myfile' "    def __init__(self, **kw):" _n
    file write `myfile' "        dict.__init__(self, kw)" _n
    file write `myfile' "        self.__dict__  = self" _n _n _n
    file write `myfile' "results = Bunch("
    foreach mat of local namelist {
        file write `myfile' "`mat'=`mat', "
    }
    file write `myfile' ")" _n _n
file close `myfile'
end

capture program drop mkarray
program define mkarray

    args mat myfile fmt
    local nrows = rowsof(`mat')
    local ncols = colsof(`mat')
    local i 1
    local j 1
    file write `myfile' "`mat' = np.array(["
    local justifyn = length("`mat' = np.array([")
    forvalues i=1/`nrows' {
        forvalues j = 1/`ncols' {
            if `i' > 1 | `j' > 1 { // then we need to indent
                forvalues k=1/`justifyn' {
                    file write `myfile' " "
                }
            }
            if `i' < `nrows' | `j' < `ncols' {
                file write `myfile' `fmt' (`mat'[`i',`j']) ", " _n
            }
            else {
                file write `myfile' `fmt' (`mat'[`i',`j'])
            }
        }
    }

    if `nrows' == 1 | `ncols' == 1 {
        file write `myfile' "])" _n _n
    }
    else {
        file write `myfile' "]).reshape(`nrows',`ncols')" _n _n
    }
end

capture program drop mat2array
program define mat2array
    version 11.2
    syntax namelist(min=1), SAVing(str) [ Format(str) APPend REPlace ]
    if "`format'"=="" local format "%16.0g"
    local saving: subinstr local saving "." ".", count(local ext)
    if !`ext' local saving "`saving'.py"
    tempname myfile
    file open `myfile' using "`saving'", write text `append' `replace'
    file write `myfile' "import numpy as np" _n _n

    foreach mat of local namelist {
        mkarray `mat' `myfile' `format'
    }
file close `myfile'
end

capture program drop mkarray
program define mkarray

    args mat myfile fmt
    local nrows = rowsof(`mat')
    local ncols = colsof(`mat')
    local i 1
    local j 1
    file write `myfile' "`mat' = np.array(["
    local justifyn = length("`mat' = np.array([")
    forvalues i=1/`nrows' {
        forvalues j = 1/`ncols' {
            if `i' > 1 | `j' > 1 { // then we need to indent
                forvalues k=1/`justifyn' {
                    file write `myfile' " "
                }
            }
            if `i' < `nrows' | `j' < `ncols' {
                file write `myfile' `fmt' (`mat'[`i',`j']) ", " _n
            }
            else {
                file write `myfile' `fmt' (`mat'[`i',`j'])
            }
        }
    }

    if `nrows' == 1 | `ncols' == 1 {
        file write `myfile' "])" _n _n
    }
    else {
        file write `myfile' "]).reshape(`nrows',`ncols')" _n _n
    }
end