diff --git a/phylogenetic/build-configs/ci/config.yaml b/phylogenetic/build-configs/ci/config.yaml
index 6f90bd9..7010b1f 100644
--- a/phylogenetic/build-configs/ci/config.yaml
+++ b/phylogenetic/build-configs/ci/config.yaml
@@ -77,6 +77,10 @@ clock_rate: 5.7e-5
 clock_std_dev: 2e-5
 divergence_units: "mutations"
 
+traits:
+  columns: ""
+  sampling_bias_correction: 3
+
 ## recency
 recency: true
 
diff --git a/phylogenetic/defaults/hmpxv1/config.yaml b/phylogenetic/defaults/hmpxv1/config.yaml
index 658b924..4cd55da 100644
--- a/phylogenetic/defaults/hmpxv1/config.yaml
+++ b/phylogenetic/defaults/hmpxv1/config.yaml
@@ -74,6 +74,10 @@ clock_rate: 5.7e-5
 clock_std_dev: 2e-5
 divergence_units: "mutations"
 
+traits:
+  columns: ""
+  sampling_bias_correction: 3
+
 ## recency
 recency: true
 
diff --git a/phylogenetic/defaults/hmpxv1_big/config.yaml b/phylogenetic/defaults/hmpxv1_big/config.yaml
index e2b63d9..1c9cf79 100644
--- a/phylogenetic/defaults/hmpxv1_big/config.yaml
+++ b/phylogenetic/defaults/hmpxv1_big/config.yaml
@@ -50,6 +50,10 @@ clock_rate: 5.7e-5
 clock_std_dev: 2e-5
 divergence_units: "mutations"
 
+traits:
+  columns: ""
+  sampling_bias_correction: 3
+
 ## recency
 recency: true
 
diff --git a/phylogenetic/defaults/mpxv/config.yaml b/phylogenetic/defaults/mpxv/config.yaml
index 8d2e39f..6bd9678 100644
--- a/phylogenetic/defaults/mpxv/config.yaml
+++ b/phylogenetic/defaults/mpxv/config.yaml
@@ -70,6 +70,10 @@ clock_rate: 3e-6
 clock_std_dev: 6e-6
 divergence_units: "mutations-per-site"
 
+traits:
+  columns: ""
+  sampling_bias_correction: 3
+
 ## recency
 recency: true
 
diff --git a/phylogenetic/rules/annotate_phylogeny.smk b/phylogenetic/rules/annotate_phylogeny.smk
index 248ea10..6a15627 100644
--- a/phylogenetic/rules/annotate_phylogeny.smk
+++ b/phylogenetic/rules/annotate_phylogeny.smk
@@ -72,8 +72,8 @@ rule traits:
     output:
         node_data=build_dir + "/{build_name}/traits.json",
     params:
-        columns="country",
-        sampling_bias_correction=3,
+        columns=config["traits"]["columns"],
+        sampling_bias_correction=config["traits"]["sampling_bias_correction"],
         strain_id=config["strain_id_field"],
     shell:
         """
diff --git a/phylogenetic/rules/export.smk b/phylogenetic/rules/export.smk
index 124b07a..1501eb2 100644
--- a/phylogenetic/rules/export.smk
+++ b/phylogenetic/rules/export.smk
@@ -32,9 +32,9 @@ OUTPUTS:
 
 rule remove_time:
     input:
-        "results/{build_name}/branch_lengths.json",
+        build_dir + "/{build_name}/branch_lengths.json",
     output:
-        "results/{build_name}/branch_lengths_no_time.json",
+        build_dir + "/{build_name}/branch_lengths_no_time.json",
     shell:
         """
         python3 scripts/remove_timeinfo.py --input-node-data {input} --output-node-data {output}
@@ -66,11 +66,15 @@ rule export:
         tree=build_dir + "/{build_name}/tree.nwk",
         metadata=build_dir + "/{build_name}/metadata.tsv",
         branch_lengths=(
-            "results/{build_name}/branch_lengths.json"
+            build_dir + "/{build_name}/branch_lengths.json"
             if config.get("timetree", False)
-            else "results/{build_name}/branch_lengths_no_time.json"
+            else build_dir + "/{build_name}/branch_lengths_no_time.json"
+        ),
+        traits=(
+            build_dir + "/{build_name}/traits.json"
+            if config.get("traits", {}).get("columns", False)
+            else []
         ),
-        traits=build_dir + "/{build_name}/traits.json",
         nt_muts=build_dir + "/{build_name}/nt_muts.json",
         aa_muts=build_dir + "/{build_name}/aa_muts.json",
         clades=build_dir + "/{build_name}/clades.json",
@@ -95,7 +99,7 @@ rule export:
             --tree {input.tree} \
             --metadata {input.metadata} \
             --metadata-id-columns {params.strain_id} \
-            --node-data {input.branch_lengths} {input.nt_muts} {input.aa_muts} {input.mutation_context} {input.clades} {input.recency}\
+            --node-data {input.branch_lengths} {input.traits} {input.nt_muts} {input.aa_muts} {input.mutation_context} {input.clades} {input.recency}\
             --colors {input.colors} \
             --lat-longs {input.lat_longs} \
             --description {input.description} \