Refactor build-states-to-change

src/clj/witan/send/model.clj

@@ -52,7 +52,8 @@
     :witan/version "1.0.0"
     :witan/type :function
     :witan/fn :send/prepare-send-inputs
-    :witan/params {:modify-transition-by 1
+    :witan/params {:which-transitions? nil
+                   :modify-transition-by 1
                    :splice-ncy nil
                    :filter-transitions-from nil}}
    {:witan/name :run-send-model

src/clj/witan/send/schemas.clj

@@ -202,3 +202,6 @@
    :setting-cost-lookup SettingCostLookup
    :valid-setting-academic-years ValidSettingAcademicYears
    :transition-matrix TransitionCounts})
+
+(def transition-type
+  (s/maybe [s/Str]))

src/clj/witan/send/send.clj

@@ -197,8 +197,20 @@
            {} a)
    b))
 
-(defn generate-transition-key [cy ay need setting]
-  (vector cy ay :NONSEND (states/state need setting)))
+(defn generate-transition-key [{:keys [transition-type cy ay need setting move-state]}]
+  (when (not= move-state (states/state need setting))
+    (case transition-type
+      "joiners"
+      (vector cy ay :NONSEND (states/state need setting))
+
+      "leavers"
+      (vector cy ay (states/state need setting) :NONSEND)
+
+      "movers-to"
+      (vector cy ay move-state (states/state need setting))
+
+      "movers-from"
+      (vector cy ay (states/state need setting) move-state))))
 
 (defn update-ifelse-assoc [m k arithmetic-fn v]
   (if (contains? m k)
@@ -260,28 +272,28 @@
               :mover-beta-params (p/beta-params-movers valid-states valid-transitions transition-matrix)
               :mover-state-alphas  (p/alpha-params-movers valid-states valid-transitions transition-matrix)}))))
 
-(defn build-states-to-change [settings-to-change valid-needs ages years]
-  (if (= :nil (-> settings-to-change
-                  ds/row-maps
-                  first
-                  :setting-2))
-    (let [states (->> settings-to-change
-                      ds/row-maps
-                      (map #(vector (:setting-1 %))))]
-      (vec (mapcat (fn [year]
-                     (mapcat (fn [age]
-                               (mapcat (fn [need]
-                                         (map (fn [setting] (vector (generate-transition-key year age need (first setting))))
-                                              states)) valid-needs)) ages)) years)))
-    (let [state-pairs (->> settings-to-change
-                           ds/row-maps
-                           (map #(vector (:setting-1 %) (:setting-2 %))))]
-      (vec (mapcat (fn [year]
-                     (mapcat (fn [age]
-                               (mapcat (fn [need]
-                                         (map (fn [setting] (vector (generate-transition-key year age need (first setting))
-                                                                    (generate-transition-key year age need (second setting))))
-                                              state-pairs)) valid-needs)) ages)) years)))))
+(defn build-states-to-change [input valid-needs valid-settings ages years transition-type]
+  (let [to-maps (ds/row-maps input)
+        settings-to-change (if (= :nil (-> to-maps
+                                           first
+                                           :setting-2))
+                             (map #(vector (:setting-1 %)) to-maps)
+                             (map #(vector (:setting-1 %) (:setting-2 %)) to-maps))]
+    (->> (for [year years
+               age ages
+               need valid-needs
+               setting valid-settings
+               setting-to-change settings-to-change]
+           (let [keys {:transition-type transition-type :cy year :ay age
+                       :need need :move-state (states/state need setting)}]
+             (if (= :nil (-> to-maps
+                             first
+                             :setting-2))
+               (vector (generate-transition-key (merge keys {:setting (first setting-to-change)})))
+               (vector (generate-transition-key (merge keys {:setting (first setting-to-change)}))
+                       (generate-transition-key (merge keys {:setting (second setting-to-change)}))))))
+         (remove #(nil? (first %)))
+         distinct)))
 
 (defworkflowfn prepare-send-inputs-1-0-0
   "Outputs the population for the last year of historic data, with one
@@ -294,7 +306,8 @@
                         :transition-matrix sc/TransitionCounts
                         :setting-cost sc/NeedSettingCost
                         :valid-setting-academic-years sc/ValidSettingAcademicYears}
-   :witan/param-schema {:modify-transition-by s/Num
+   :witan/param-schema {:which-transitions? sc/transition-type
+                        :modify-transition-by s/Num
                         :splice-ncy (s/maybe sc/AcademicYear)
                         :filter-transitions-from (s/maybe [sc/CalendarYear])}
    :witan/output-schema {:standard-projection sc/projection-map
@@ -303,7 +316,7 @@
                          :settings-to-change sc/SettingsToChange}}
   [{:keys [settings-to-change initial-send-population transition-matrix population
            setting-cost valid-setting-academic-years]}
-   {:keys [modify-transition-by splice-ncy filter-transitions-from]}]
+   {:keys [which-transitions? modify-transition-by splice-ncy filter-transitions-from]}]
   (let [original-transitions transition-matrix
         ages (distinct (map :academic-year (ds/row-maps population)))
         years (distinct (map :calendar-year (ds/row-maps population)))
@@ -314,7 +327,7 @@
         valid-states (states/calculate-valid-states-from-setting-academic-years initialise-validation)
         valid-year-settings (states/calculate-valid-year-settings-from-setting-academic-years initialise-validation)
         states-to-change (when (not= 1 modify-transition-by)
-                           (build-states-to-change settings-to-change valid-needs ages years))
+                           (mapcat (fn [transition-type] (build-states-to-change settings-to-change valid-needs valid-settings ages years transition-type)) which-transitions?))
         transition-matrix (ds/row-maps transition-matrix)
         modified-transition-matrix (when (not= 1 modify-transition-by)
                                      (let [convert (-> transition-matrix

test/witan/send/acceptance/workspace_test.clj

@@ -55,16 +55,21 @@
 
   :transitions-file - csv containing list of settings to modify by :transition-modifier (optional)
 
+  :which-transitions? - vector of either one or more transition type strings i.e. \"joiners\", \"movers-to\",
+  \"movers-from\" or \"leavers\", used in conjunction with :transition-modifier and :transitions-file (optional)
+
   :modify-transitions-from - set a year to start modifying transitions from when :transition-modifier & :transitions-file are set (optional)
 
   :filter-transitions-from - sets a year or year range as a vector to filter historic transitions by for :splice-ncy (optional)
 
   :splice-ncy - sets a national curriculum year to ignore transitions of prior to :filter-transitions-from year (optional)"
-  [{:keys [iterations output? transition-modifier transitions-file modify-transitions-from filter-transitions-from splice-ncy]}]
+  [{:keys [iterations output? transition-modifier transitions-file which-transitions?
+           modify-transitions-from filter-transitions-from splice-ncy]}]
   (report/reset-send-report)
   (report/info "Input Data: " (report/bold (str/replace (str/join "" (drop-last inputs-path)) #"/" " ")))
   (report/info "Number of iterations: " (report/bold iterations))
   (report/info "Output charts produced: " (report/bold output?))
+  (report/info "Modifying " (report/bold (if (nil? which-transitions?) "None" (str/join ", " which-transitions?))))
   (report/info "Transitions modifier: " (report/bold (if (nil? transition-modifier) "None" transition-modifier)))
   (report/info "Transitions file: " (report/bold (if (nil? transitions-file) "None" transitions-file)))
   (report/info "Modify transitions from: " (report/bold (if (nil? modify-transitions-from) "None" modify-transitions-from)))
@@ -80,7 +85,9 @@
                                                (map #(assoc-in % [:witan/params :output] output?)))
                             prep-catalog2 (if (nil? transition-modifier)
                                             prep-catalog1
-                                            (map #(assoc-in % [:witan/params :modify-transition-by] transition-modifier) prep-catalog1))
+                                            (->> prep-catalog1
+                                                 (map #(assoc-in % [:witan/params :modify-transition-by] transition-modifier))
+                                                 (map #(assoc-in % [:witan/params :which-transitions?] which-transitions?))))
                             prep-catalog3 (if (nil? modify-transitions-from)
                                             prep-catalog2
                                             (map #(assoc-in % [:witan/params :modify-transitions-from] modify-transitions-from) prep-catalog2))]

test/witan/send/send_test.clj

@@ -199,3 +199,43 @@
     (is (= 2 (:foo (update-ifelse-assoc {:foo 1 :bar 2} :foo + 1)))))
   (testing "if key not present, insert key with val"
     (is (= 1 (:foo (update-ifelse-assoc {:baz 1 :bar 2} :foo + 1))))))
+
+(deftest generate-transition-key-test
+  (testing "generate joiner state transition"
+    (is (= [2013 5 :NONSEND :CI-MSSOB]
+           (generate-transition-key {:transition-type "joiners" :cy 2013 :ay 5
+                                     :need :CI :move-state :CI-MSSIB
+                                     :setting :MSSOB}))))
+  (testing "generate leaver state transition"
+    (is (= [2013 5 :CI-MSSOB :NONSEND]
+           (generate-transition-key {:transition-type "leavers" :cy 2013 :ay 5
+                                     :need :CI :move-state :CI-MSSIB
+                                     :setting :MSSOB}))))
+  (testing "generate joiner state transition"
+    (is (= [2013 5 :CI-MSSIB :CI-MSSOB]
+           (generate-transition-key {:transition-type "movers-to" :cy 2013 :ay 5
+                                     :need :CI :move-state :CI-MSSIB
+                                     :setting :MSSOB}))))
+  (testing "generate joiner state transition"
+    (is (= [2013 5 :CI-MSSOB :CI-MSSIB]
+           (generate-transition-key {:transition-type "movers-from" :cy 2013 :ay 5
+                                     :need :CI :move-state :CI-MSSIB
+                                     :setting :MSSOB})))))
+
+(deftest build-states-to-change-test
+  (let [input (ds/dataset [{:setting-1 :MMSOB :setting-2 :MMSIB}])]
+    (testing "move MSSOB joiners to MSSIB"
+      (is (= [[2013 5 :NONSEND :CI-MMSOB] [2013 5 :NONSEND :CI-MMSIB]]
+             (first (build-states-to-change input [:CI] [:MSSOB] [5] [2013] "joiners")))))
+    (testing "move MSSOB leavers to MSSIB"
+      (is (= [[2013 5 :CI-MMSOB :NONSEND] [2013 5 :CI-MMSIB :NONSEND]]
+             (first (build-states-to-change input [:CI] [:MSSOB] [5] [2013] "leavers")))))
+    (testing "move MSSOB movers-to to MSSIB"
+      (is (= [[2013 5 :CI-MSSOB :CI-MMSOB] [2013 5 :CI-MSSOB :CI-MMSIB]]
+             (first (build-states-to-change input [:CI] [:MSSOB] [5] [2013] "movers-to")))))
+    (testing "move MSSOB movers-from to MSSIB"
+      (is (= [[2013 5 :CI-MMSOB :CI-MSSOB] [2013 5 :CI-MMSIB :CI-MSSOB]]
+             (first (build-states-to-change input [:CI] [:MSSOB] [5] [2013] "movers-from")))))
+    (testing "return nil when move to and from state the same"
+      (is (= nil
+             (first (build-states-to-change input [:CI] [:MMSOB] [5] [2013] "movers-to")))))))